nodes.py

"""
pydtnsim node module.

Implement non clustering routing algorithms here
- Epidemic
- Flooding
- PRoPHET
- PRoPHET single copy
"""

__all__ = [
    "EpidemicNode",
    "FloodingNode",
    "ProphetNode",
    "ProphetNodeSingle",
]
from collections import defaultdict
from pydtnsim import Node


class EpidemicNode(Node):
    """Node which forwards epidemically."""

    def __init__(self, **options):
        """Create an epidemic node."""
        super().__init__(**options)
        self.sent = defaultdict(set)

    def forward(self, packet):
        """
        Forward based on the epidemic heuristic.

        Arguments:
        packet -- packet to be forwarded.

        Epidemic Heuristic:
        - Forward the packet to all neighbours that I have not forwarded
          to so far.

        """
        forward = {
            neighbour: "epidemic"
            for neighbour in self.neighbours
            if neighbour not in self.sent[packet]
        }

        return forward

    def send_success(self, packet, target):
        """
        Call when a send succeeds.

        Arguments:
        packet -- packet that was sent successfully.
        target -- the node that packet was sent to.

        Adds target to packet tracking.
        Also overrides default of deleting from buffer.
        """
        self.sent[packet].add(target)

    def packet_expiry(self, packet):
        """
        Call when a packet expires.

        Removes packet from tracking.
        """
        super().packet_expiry(packet)
        if packet in self.sent:
            del self.sent[packet]


class FloodingNode(Node):
    """Node which forwards through flooding."""

    def forward(self, packet):
        """Forward to all neighbour nodes."""
        forward = {neighbour: "flood" for neighbour in self.neighbours}
        return forward

    def send_success(self, packet, target):
        """
        Call when send suecceeds.

        Do nothing, overrides default of removing from buffer.
        """


class ProphetNode(EpidemicNode):
    """Node which implements PRoPHET, an improvement to epidemic."""

    def __init__(self, **options):
        """Create a PRoPHET node."""
        super().__init__(**options)
        self._delivery_probability = defaultdict(float)
        self.initialization = 0.75
        self.ageing = 0.98
        self.transitivity = 0.25

    def forward(self, packet):
        """Forward packet with the PRoPHET heuristic."""
        # direct forwarding
        forward = Node.forward(self, packet)
        if forward and packet.destination not in self.sent[packet]:
            return forward

        # transmit to node with the highest delivery probability,
        # if I have not already sent it to them
        target = self
        dest = packet.destination.name
        for node in self.neighbours:
            if node.delivery_probability_to(
                dest
            ) > target.delivery_probability_to(dest):
                target = node
        if target != self and target not in self.sent[packet]:
            return {target: "PRoPHET"}

        return {}

    def delivery_probability_to(self, dest):
        """Return delivery probability to given destination."""
        if dest not in self.delivery_probability:
            return 0
        return self.delivery_probability[dest]

    def tick(self):
        """Apply aging constant every tick."""
        # pylint:disable=duplicate-code
        tick = self.options["tick_rate"]

        while True:
            start = self.network.now

            yield from self.forward_all()
            delay = tick - (self.network.now - start) % tick

            # apply aging constant
            self._delivery_probability = {
                node: prob * self.ageing
                for node, prob in self.delivery_probability.items()
            }

            yield self.network.env.timeout(delay)

    @property
    def delivery_probability(self):
        """Return delivery probability to all known destinations."""
        return dict(self._delivery_probability)

    def join(self, node):
        """Update PRoPHET delivery probabilities on contact with other node."""
        super().join(node)

        index = node.name  # index by name to avoid pickling generator error
        # initial delivery probability
        if index not in self.delivery_probability:
            self._delivery_probability[index] = 0

        old = self.delivery_probability[index]
        self._delivery_probability[index] = (
            old + (1 - old) * self.initialization
        )

        # transitive delivery probability
        for other, probability in node.delivery_probability.items():
            if other == self.name:
                pass
            if other not in self.delivery_probability:
                self._delivery_probability[other] = 0
            old = self.delivery_probability[other]
            self._delivery_probability[other] = (
                old
                + (1 - old)
                * self.delivery_probability[index]
                * probability
                * self.transitivity
            )


class ProphetNodeSingle(ProphetNode):
    """PRoPHET node with single copy routing."""

    def send_success(self, packet, target):
        """
        Call when send suecceeds.

        overrides PRoPHET default of keeping in buffer.
        """
        self.buffer.remove(packet)