# Optimising Graphed Rewards

Relates to this and this.

Simply put; I want to be able to distribute a total reward among network, where the distribution diminishes according to the depth of the network but the total sum of distributed rewards equals the initial total reward. As well as this I want to be able to support graph structures of the form:

1. Self referencing (Node0 -> Node0)
2. Linear Chains (Node0 -> Node1 -> Node2 -> Node3)
3. Circular Chains (Node0 -> Node1 -> Node2 -> Node0)

We will always have a reference start node which should always get a fixed amount, then we should distribute the remainder up-stream. I have tried to think in terms of invariants with all the pre-conditions and correctness over performance.

As part of a review I would like guidance on:

1. How the solution becomes as close to perfect with regards to rounding reward amounts.
2. Performance considerations; how do we make this process more efficient?
3. Is there anything else I can use to inform a tighter bound?
4. I'm sure NetworkX offers various built-in methods which I am not using, so guidance on that would be great

This is not an exhaustive list, if there are other issues, please include in a review.

### Code

import networkx as nx
from collections import deque

def is_circular(graph: nx.DiGraph, start_node: str) -> bool:
"""
Determines if there is a cycle in a directed graph that includes a specific start node.

This function traverses the given directed graph to find all simple cycles and checks
if any of these cycles include the specified start node. A simple cycle is a cycle in a
graph that does not repeat any nodes except the starting and ending node.

Parameters:
- graph (nx.DiGraph): The directed graph to analyze, represented as a NetworkX DiGraph object.
- start_node (str): The identifier of the node from which to check for cyclicity. This is
expected to be a string that matches one of the node identifiers in the graph.

Returns:
- bool: True if there is at least one cycle in the graph that includes the start node,
False otherwise.
"""
# If any cycle includes the start node, return True
cycles = nx.simple_cycles(graph)
for cycle in cycles:
if start_node in cycle:
return True
return False

def distribute_rewards(
graph: nx.DiGraph,
start_node: str,
total_reward: float,
start_node_ratio: float,
upward_referral_ratio: float,
loop_reward_ratio: float,
max_referral_levels: int
) -> None:
"""
Distributes rewards across a referral network represented as a directed graph.
The distribution considers direct, upward referrals, and rewards for cycles.

Parameters:
- graph (nx.DiGraph): Directed graph representing the referral network.
- start_node (str): The node from which distribution starts.
- total_reward (float): The total amount of reward to be distributed.
- start_node_ratio (float): The fraction of total reward allocated to the start node.
- upward_referral_ratio (float): The fraction of remaining reward distributed to upward referrals.
- loop_reward_ratio (float): The fraction of total reward allocated when a loop is detected.
- max_referral_levels (int): The maximum number of referral levels to distribute rewards through.

The function modifies the graph in-place by setting a 'reward' attribute for each node.
"""

assert 0 <= start_node_ratio <= 1, "Start node ratio must be between 0 and 1"
assert 0 <= upward_referral_ratio <= 1, "Upward referral ratio must be between 0 and 1"
assert 0 <= loop_reward_ratio <= 1, "Loop reward ratio must be between 0 and 1"

nx.set_node_attributes(graph, 0.0, 'reward')

# Check if the graph has a self-referencing node
if graph.has_edge(start_node, start_node):
# Directly assign the designated portion of the total reward to the self-referencing node
# and bypass the distribution logic meant for non-self-referencing scenarios
self_referral_reward = total_reward * loop_reward_ratio
graph.nodes[start_node]['reward'] = self_referral_reward
return

visited = set()
start_node_reward = total_reward * (loop_reward_ratio if is_circular(graph, start_node) else start_node_ratio)
graph.nodes[start_node]['reward'] = start_node_reward
remaining_reward = total_reward - start_node_reward

distributed_rewards = {start_node: start_node_reward}
queue = deque([(start_node, 0)])

max_iterations = len(graph.nodes()) * max_referral_levels
iteration_count = 0
while queue and iteration_count < max_iterations:
iteration_count += 1
current_node, level = queue.popleft()
if 0 < level <= max_referral_levels:
reward_for_node = (upward_referral_ratio * remaining_reward) / (2 ** (level - 1))
distributed_rewards[current_node] = reward_for_node
graph.nodes[current_node]['reward'] += reward_for_node  # Apply reward

# Adjusting the node exploration based on context: circular or linear
next_nodes = (
graph.successors(current_node)
if is_circular(graph, start_node)
else graph.predecessors(current_node)
)
for next_node in next_nodes:
if next_node not in visited:
queue.append((next_node, level + 1))

# Post-distribution correction if necessary
sum_distributed_rewards = sum(distributed_rewards.values())
if sum_distributed_rewards != start_node_reward:
adjustment_ratio = (total_reward - start_node_reward) / (max(sum_distributed_rewards - start_node_reward, 1))
for node, reward in distributed_rewards.items():
# Avoid re-adjusting the start node
if node != start_node:
corrected_reward = reward * adjustment_ratio
graph.nodes[node]['reward'] = corrected_reward

final_sum_rewards = sum(nx.get_node_attributes(graph, 'reward').values())
assert abs(final_sum_rewards - total_reward) < 1e-6, "The sum of node rewards does not match the total reward."


## Tests

import pytest
import networkx as nx
from typing import TypedDict, Unpack, Protocol

from rewarder import distribute_rewards, is_circular

class RewardKwargs(TypedDict):
total_reward: float
start_node_ratio: float
upward_referral_ratio: float
loop_reward_ratio: float
max_referral_levels: int

class SimpleGraphMaker(Protocol):
def __call__(self, edges: list[tuple[str, ...]]) -> nx.DiGraph:
...

class RewardGraphMaker(Protocol):
def __call__(self, edges: list[tuple[str, ...]], start_node: str, **kwargs: Unpack[RewardKwargs]) -> nx.DiGraph:
...

@pytest.fixture
def create_graph() -> SimpleGraphMaker:

def _create_graph(edges: list[tuple[str, ...]]) -> nx.DiGraph:
graph = nx.DiGraph()
for edge in edges:
return graph

return _create_graph

@pytest.fixture
def create_reward_graph() -> RewardGraphMaker:

def _create_graph(edges: list[tuple[str, ...]], start_node: str, **kwargs: Unpack[RewardKwargs]) -> nx.DiGraph:
graph = nx.DiGraph()
for edge in edges:

distribute_rewards(
graph,
start_node,
**kwargs
)
return graph

return _create_graph

def assert_rewards(graph: nx.DiGraph, expected_rewards: dict[str, float]) -> None:
for node, expected_reward in expected_rewards.items():
assert graph.nodes[node]['reward'] == pytest.approx(expected_reward, 0.01), f"Reward mismatch for {node}"

def test_is_circular_with_circular_graph(create_graph: SimpleGraphMaker) -> None:
graph = create_graph(
edges=[
('Node0', 'Node1'),
('Node1', 'Node2'),
('Node2', 'Node0')
],
)
assert is_circular(graph, 'Node0') is True

def test_is_circular_with_linear_graph(create_graph: SimpleGraphMaker) -> None:
graph = create_graph(
edges=[
('Node0', 'Node1'),
('Node1', 'Node2'),
('Node2', 'Node3')
],
)
assert is_circular(graph, 'Node0') is False

def test_is_circular_with_self_loop(create_graph: SimpleGraphMaker) -> None:
graph = create_graph(
edges=[
('Node0', 'Node0')
]
)
assert is_circular(graph, 'Node0') is True

def test_is_circular_with_empty_graph(create_graph: SimpleGraphMaker) -> None:
graph = create_graph(
edges=[]
)
assert is_circular(graph, 'Node0') is False

def test_is_circular_nonexistent_start_node(create_graph: SimpleGraphMaker) -> None:
graph = create_graph(
edges=[
('Node1', 'Node2'),
('Node2', 'Node3')
]
)
assert is_circular(graph, 'Node0') is False

def test_is_circular_complex_graph(create_graph: SimpleGraphMaker) -> None:
graph = create_graph(
edges=[
('Node0', 'Node1'),
('Node1', 'Node2'),
('Node2', 'Node3'),
('Node3', 'Node1')
]
)
assert is_circular(graph, 'Node1') is True
assert is_circular(graph, 'Node0') is False  # Node0 is not part of the cycle

# Self Referencing Graphs of the form Node0 --> Node0

def test_self_referencing(create_reward_graph: RewardGraphMaker) -> None:
graph = create_reward_graph(
edges=[
('Node0', 'Node0'),
],
start_node='Node0',
**RewardKwargs(
total_reward=1333,
start_node_ratio=0.85,
upward_referral_ratio=0.15,
loop_reward_ratio=0.85,
max_referral_levels=10
)
)
expected_rewards = {
'Node0': 1133.05
}
assert_rewards(graph, expected_rewards)

# Circular Graphs are of the form Node0 --> Node1 --> Node2 --> Node0

def test_circular(create_reward_graph: RewardGraphMaker) -> None:
graph = create_reward_graph(
edges=[
('Node0', 'Node1'),
('Node1', 'Node2'),
('Node2', 'Node0')
],
start_node='Node0',
**RewardKwargs(
total_reward=1333,
start_node_ratio=0.85,
upward_referral_ratio=0.15,
loop_reward_ratio=0.85,
max_referral_levels=10
)
)
expected_rewards = {
'Node0': 1133.05,
'Node1': 133.3,
'Node2': 66.65,
}
assert_rewards(graph, expected_rewards)

# Chained Graphs are of the form Node0 --> Node1 --> Node2 --> Node3 --> NodeN

def test_max_level_chained(create_reward_graph: RewardGraphMaker) -> None:
"""10 levels deep (most likely the maximum we will go), Node9 triggers the reward."""
graph = create_reward_graph(
edges=[
('Node0', 'Node1'),
('Node1', 'Node2'),
('Node2', 'Node3'),
('Node3', 'Node4'),
('Node4', 'Node5'),
('Node5', 'Node6'),
('Node6', 'Node7'),
('Node7', 'Node8'),
('Node8', 'Node9')
],
start_node='Node9',
**RewardKwargs(
total_reward=1333,
start_node_ratio=0.85,
upward_referral_ratio=0.15,
loop_reward_ratio=0.85,
max_referral_levels=10
)
)
expected_rewards = {
'Node9': 1133.05,
'Node8': 100.17,
'Node7': 50.08,
'Node6': 25.04,
'Node5': 12.52,
'Node4': 6.26,
'Node3': 3.13,
'Node2': 1.56,
'Node1': 0.782,
'Node0': 0.391
}
assert_rewards(graph, expected_rewards)

def test_custom_level_chained(create_reward_graph: RewardGraphMaker) -> None:
"""max_level=10, but tree is only 5 deep. Node5 triggers the reward."""
graph = create_reward_graph(
edges=[
('Node0', 'Node1'),
('Node1', 'Node2'),
('Node2', 'Node3'),
('Node3', 'Node4'),
('Node4', 'Node5')
],
start_node='Node5',
**RewardKwargs(
total_reward=1333,
start_node_ratio=0.85,
upward_referral_ratio=0.15,
loop_reward_ratio=0.85,
max_referral_levels=10
)
)
expected_rewards = {
'Node5': 1133.05,
'Node4': 103.19,
'Node3': 51.59,
'Node2': 25.79,
'Node1': 12.89,
'Node0': 6.44
}
assert_rewards(graph, expected_rewards)


is_circular can be rephrased as

return any(
start_node in cycle
for cycle in nx.simple_cycles(graph)
)


Try to redesign the API for distribute_rewards so that it doesn't mutate its arguments. This might require that you break up the function into subroutines.

The asserts are useful - so much so that they should be kept even over varying Python arguments, so just convert them to conditional raise statements.

Convert this loop:

    iteration_count = 0
while queue and iteration_count < max_iterations:
iteration_count += 1


so that you don't manually increment, and instead iterate over an itertools.count.

(max(sum_distributed_rewards - start_node_reward, 1)) does not need outer parens.

Seemingly reasonable types. If you have a mypy.ini, be sure to try it with strict mode on. If too many things break and it's not possible to fix them, work backwards, adding exceptions where needed.

Good tests!

The business case is complicated and I can't intelligently speak to it, but if your tests pass and they're based on reality, that's a good start.