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