def test_repr_3(self):
"""
Base case: Nodes and Marbles given.
"""
marble_pos = torch.tensor([1.])
marble_vel = torch.tensor([2.])
marble_acc = torch.tensor([3.])
marble = Marble(marble_pos, marble_vel, marble_acc, 8, None, None, 0.4,
0.5, 0.6, 0.7)
marble_2 = marble.copy()
node = Node(ZERO, ZERO, ZERO, 0, None, 0, 0, 0, 0)
node_2 = Node(torch.tensor([11.]), torch.tensor([12.]),
torch.tensor([13.]), 14, None, 0, 0, 0, 0)
expected = f"NenwinModel({repr(set([node, node_2]))}," \
+ f"{repr(set([marble, marble_2]))})"
result = repr(NenwinModel([node, node_2], [marble, marble_2]))
self.assertEqual(expected, result)
def test_reset_removed_added_marbles(self):
"""
NenwinModel.reset() should remove all Marbles
added with NenwinModel.add_marbles()
"""
original_marble = Marble(ZERO,
ZERO,
ZERO,
10,
NewtonianGravity(),
datum=None)
added_marble = original_marble.copy()
assert added_marble is not original_marble
model = NenwinModel([], [original_marble])
model.add_marbles([added_marble])
model.reset()
self.assertIn(original_marble, model.marbles)
self.assertIn(original_marble, model._NenwinModel__all_particles)
self.assertNotIn(added_marble, model.marbles)
self.assertNotIn(added_marble, model._NenwinModel__all_particles)