gears.py

#!/usr/bin/python

# First some classes to implement graph theory. These aren't supposed to know
# anything about Physics.

class Node(object):
    """Nodes are told about the edges that exist as they are created, so 
    they know about each other. This means that a set of connected Nodes can be
    referred to by any one of them; none are special."""

    def __init__(self):
        self.edges = set()

    def edge_added(self, edge):
        self.edges.add(edge)

    def edge_removed(self, edge):
        self.edges.remove(edge)

    def connected_nodes(self):
        """Return a frozenset of other nodes that this node is connected to."""
        return frozenset(edge.other_node(self) for edge in self.edges)

    def traverse_breadth_first(self, encountered=None):
        """Generate all connected nodes in the graph breadth-first starting
        with self. The encountered parameter is used internally for
        recursion."""
        if encountered is None:
            encountered = set([self])

        yield self

        all_connected_nodes = self.connected_nodes()
        nodes_to_recurse = all_connected_nodes.difference(encountered)

        encountered.update(nodes_to_recurse)
        for node in nodes_to_recurse:
            yield node

        for node in nodes_to_recurse:
            node.traverse_breadth_first(encountered)

class Edge(object):
    """Edges are used to manage the graph, but do not exist within the graph
    until they are "activated". When an Edge is activated it notifies the
    relevant Nodes that they are connected with it so that they are able to
    discover and traverse the entire connected part of the graph."""

    # TODO: how would we find disjointed sections in order to run the physics
    # on each separately?

    def __init__(self, a, b):
        self.nodes = (a, b)

    def other_node(self, node):
        assert(node in self.nodes)
        return self.nodes[node is not self.nodes[0]]

    def activate(self):
        for node in self.nodes:
            node.edge_added(self)

    def deactivate(self):
        for node in self.nodes:
            node.edge_removed(self)


# Physics starts here

class Link(Edge):
    """A Link connects two Physics objects that we want to use as an Edge in
    the graph. Only particular types of Physics objects are able to link, which
    they define later. Linked Physics objects have their Edges automatically
    activated."""

    def __init__(self, a, b):
        assert(a.can_link_to(b))
        assert(b.can_link_to(a))
        
        Edge.__init__(self, a, b)
        self.activate()

class PhysicsNode(Node):
    @classmethod
    def can_link_to(cls, other):
        return any(isinstance(other, t) for t in cls._can_link_to)

class Gear(PhysicsNode):
    def __init__(self, mass, radius):
        # Static properties
        self.mass = mass
        self.radius = radius

        # Dynamic properties
        self.angular_velocity = 0.0
        self.phase = 0.0 # angle in radians

    _can_link_to = [ Mesh, Hammer, Drive ] # Load is implicit by inheritence

class Mesh(PhysicsNode):
    """A Mesh represents the interface between two meshed Gears. In a static
    analysis, it has two identical forces of equal magnitude acting on the
    connected Gears in opposite directions."""

    def __init__(self, gear1, gear2):
        self.edges = (Link(self, gear1), Link(self, gear2))

    _can_link_to = [ Gear ]

class Hammer(PhysicsNode):
    """A Hammer's shaft is connected to the axle of one Gear at one end, and a
    mass at the other end. In a static analysis, it has a torque (Nm) at the
    Gear end, and a force (N) at the mass end."""

    def __init__(self, mass, length):
        # Static properties
        self.mass = mass
        self.length = length

    # Dynamic properties taken from the gear it is connected to
    @property
    def phase(self):
        return self.connected_gear().phase

    @property
    def angular_velocity(self):
        return self.connected_gear().angular_velocity

    def connected_gear(self):
        # For now, a Hammer can only be connected to one thing, which must
        # be the Gear
        assert(len(self.edges) == 1)
        return list(self.edges)[0].other_node(self)
        
    _can_link_to = [ Gear ]

class Drive(PhysicsNode):
    """A Drive supplies torque to a Gear. It does Work."""
    def __init__(self, torque):
        self.torque = torque

    _can_link_to = [ Gear ]

class Load(Drive):
    """A Load is just an inverse Drive."""
    def __init__(self, torque):
        Drive.__init__(-torque)

    # _can_link_to = [ Gear ] as defined by Drive