class Satellite(RigidBody):
d = 3 # Cartesian Embedding dimension
D = 9 #=3 euler coordinate dimension
#angular_dims = range(3)
n = 12
def __init__(self, mass=1, l=1):
self.body_graph = BodyGraph()
self.body_graph.add_extended_nd(0, m=12, moments=(1 / 2, 1 / 2,
1 / 2)) #main body
self.body_graph.add_extended_nd(1, m=3, moments=(1, 1, 1 / 8))
self.body_graph.add_joint(0,
torch.tensor([1., 0, 0]),
1,
torch.tensor([0., 0, 0]),
rotation_axis=(torch.tensor([1., 0, 0]),
torch.tensor([0, 0, 1.])))
self.body_graph.add_extended_nd(2, m=3, moments=(1, 1, 1 / 8))
self.body_graph.add_joint(0,
torch.tensor([0., -1, 0]),
2,
torch.tensor([0., 0, 0]),
rotation_axis=(torch.tensor([0., -1, 0]),
torch.tensor([0, 0, 1.])))
# self.body_graph.add_extended_nd(3,m=3,moments=(1,1,1/8))
# self.body_graph.add_joint(2,torch.tensor([1.,1,1]),3,torch.tensor([0.,0,0]),
# rotation_axis=(torch.tensor([1.,1,1])/np.sqrt(3),torch.tensor([0,0,1.])))
def sample_initial_conditions(self, N):
bodyX = torch.randn(N, 2, self.n, self.d)
return project_onto_constraints(self.body_graph, bodyX)
def potential(self, x):
""" Gravity potential """
return 0
class ChainPendulumV2(ChainPendulum):
def __init__(self, links=2, beams=False, m=1, l=1):
self.body_graph = BodyGraph() #nx.Graph()
self.arg_string = f"n{links}{'b' if beams else ''}m{m}l{l}"
beam_moments = torch.tensor([m * l * l / 12])
if beams:
self.body_graph.add_extended_nd(0, m=m, moments=beam_moments, d=1)
self.body_graph.add_joint(0, torch.tensor([l / 2]))
for i in range(1, links):
self.body_graph.add_extended_nd(i,
m=m,
moments=beam_moments,
d=1)
self.body_graph.add_joint(i - 1, torch.tensor([-l / 2]), i,
torch.tensor([l / 2]))
else:
self.body_graph.add_node(0, m=m, tether=torch.zeros(2), l=l)
for i in range(1, links):
self.body_graph.add_node(i, m=m)
self.body_graph.add_edge(i - 1, i, l=l)
class Gyroscope(RigidBody):
d=3 # Cartesian Embedding dimension
D=3 #=3 euler coordinate dimension
angular_dims = range(3)
n=4
dt=0.02
integration_time = 2
def __init__(self, mass=.1, obj='gyro'):
verts,tris = read_obj(obj+'.obj')
verts[:,2] -= verts[:,2].min() # set bottom as 0
_,com,covar = compute_moments(torch.from_numpy(verts[tris]))
print(torch.diag(torch.diag(covar).sum()*torch.eye(3)-covar),torch.diag(covar))
self.obj = (verts,tris,com.numpy())
self.body_graph = BodyGraph()
self.body_graph.add_extended_nd(0,m=mass,moments=100*torch.diag(covar))
self.body_graph.add_joint(0,-com,pos2=torch.tensor([0.,0.,0.]))
def sample_initial_conditions(self,N):
# comEulers = torch.randn(N,2,6)
# comEulers[:,1,:3]=0
# comEulers[:,0,:3]=torch.tensor([0.,0.,1.])+.3*torch.randn(3)
# comEulers[:,0,3:] = 1*torch.randn(3)
# comEulers[:,1,3:] *=0#2#.5
# comEulers[:,1,5] = 5
# bodyX = comEuler2bodyX(comEulers)
# try: return project_onto_constraints(self.body_graph,bodyX,tol=1e-5)
# except OverflowError: return self.sample_initial_conditions(N)
eulers = (torch.rand(N,2,3)-.5)*3
#eulers[:,0,1]*=.2
eulers[:,1,0]*=3
eulers[:,1,1]*=.2
eulers[:,1,2] = (torch.randint(2,size=(N,)).float()*2-1)*(torch.randn(N)+7)*1.5
return self.body2globalCoords(eulers)
def body2globalCoords(self,eulers):
""" input: (bs,2,3) output: (bs,2,4,3) """
coms = torch.zeros_like(eulers)
comEulers = torch.cat([coms,eulers],dim=-1)
bodyX = comEuler2bodyX(comEulers)
# need to offset x,v so that joint is stationary
# pos joint =
body_attachment = self.body_graph.nodes[0]['joint'][0].to(eulers.device,eulers.dtype)
ct = torch.cat([1-body_attachment.sum()[None],body_attachment])
global_coords_attachment_point = (bodyX*ct[:,None]).sum(-2,keepdims=True) #(bs,2,3)
return bodyX-global_coords_attachment_point
def global2bodyCoords(self,bodyX):
""" input: (bs,2,4,3) output: (bs,2,3)"""
eulers = bodyX2comEuler(bodyX)[...,3:] # unwrap the euler angles
eulers[:,0,:] = torch.from_numpy(np.unwrap(eulers[:,0,:].numpy(),axis=0)).to(bodyX.device,bodyX.dtype)
# print(eulers[:,0])
# assert False
return eulers
def potential(self, x):
""" Gravity potential """
return 9.81*(self.M @ x)[..., 2].sum(1)
@property
def animator(self):
return RigidAnimation