diagrams-opus-qt/pipeline_editor/calc/flow_solver.py

"""Mock flow-distribution 'calculation'.

This is a deliberately simplified stand-in for a real hydraulic/electrical
solver. It distributes flow from source nodes to sink nodes across the network
so that each edge receives a plausible signed flow value and direction, which
the animation layer then visualizes.

Approach (mock, not physically rigorous):
  1. Identify sources (supply/feed) and sinks (demand) from node properties.
  2. Build an undirected graph of nodes connected by edges.
  3. For each (source, sink) pair with a path, push the sink's demand (scaled by
     available supply) along the shortest path, accumulating signed flow per edge.
  4. Edge direction is the sign of the accumulated flow along source->target.

The result is stored on each :class:`EdgeModel` as ``flow`` (magnitude) and
``flow_direction`` (+1/-1/0).
"""
from __future__ import annotations

from collections import deque
from typing import Optional


class FlowResult:
    """Per-edge flow outcome plus network totals."""

    def __init__(self) -> None:
        self.edge_flow: dict[str, float] = {}       # signed, along source->target
        self.total_supply: float = 0.0
        self.total_demand: float = 0.0
        self.unbalanced: float = 0.0                # supply - served demand

    def direction(self, edge_id: str) -> int:
        f = self.edge_flow.get(edge_id, 0.0)
        if f > 1e-9:
            return 1
        if f < -1e-9:
            return -1
        return 0


def _node_supply(node) -> float:
    return float(node.properties.value("supply", 0.0) or 0.0)


def _node_demand(node) -> float:
    return float(node.properties.value("demand", 0.0) or 0.0)


def solve(document) -> FlowResult:
    """Distribute flow across the document's edges. Returns a :class:`FlowResult`."""
    result = FlowResult()
    nodes = list(document.nodes())
    edges = list(document.edges())
    if not edges:
        return result

    # adjacency: node_id -> list of (neighbor_id, edge_id, sign_for_source_to_target)
    adj: dict[str, list[tuple[str, str, int]]] = {n.node_id: [] for n in nodes}
    for e in edges:
        if e.source_node in adj and e.target_node in adj:
            adj[e.source_node].append((e.target_node, e.edge_id, +1))
            adj[e.target_node].append((e.source_node, e.edge_id, -1))
        result.edge_flow[e.edge_id] = 0.0

    sources = [(n.node_id, _node_supply(n)) for n in nodes if _node_supply(n) > 0]
    sinks = [(n.node_id, _node_demand(n)) for n in nodes if _node_demand(n) > 0]
    result.total_supply = sum(s for _, s in sources)
    result.total_demand = sum(d for _, d in sinks)

    if not sources or not sinks:
        # No defined supply/demand: assign a nominal unit flow so the diagram
        # still animates. Direction follows edge source->target.
        for e in edges:
            result.edge_flow[e.edge_id] = 1.0
        return result

    # Scale served demand so we never exceed available supply.
    scale = 1.0
    if result.total_demand > result.total_supply and result.total_demand > 0:
        scale = result.total_supply / result.total_demand
    result.unbalanced = result.total_supply - result.total_demand * scale

    remaining_supply = {sid: s for sid, s in sources}
    for sink_id, demand in sinks:
        served = demand * scale
        # distribute this sink's served demand across reachable sources
        while served > 1e-9:
            src = _nearest_source_with_supply(sink_id, remaining_supply, adj)
            if src is None:
                break
            path = _shortest_path(src, sink_id, adj)
            if not path:
                remaining_supply[src] = 0.0
                continue
            push = min(served, remaining_supply[src])
            _accumulate_along_path(path, adj, push, result)
            remaining_supply[src] -= push
            served -= push

    return result


def _nearest_source_with_supply(sink_id, remaining_supply, adj) -> Optional[str]:
    """BFS from sink to the closest source that still has supply."""
    seen = {sink_id}
    q = deque([sink_id])
    while q:
        cur = q.popleft()
        if cur in remaining_supply and remaining_supply[cur] > 1e-9:
            return cur
        for nb, _eid, _sign in adj.get(cur, []):
            if nb not in seen:
                seen.add(nb)
                q.append(nb)
    return None


def _shortest_path(start, goal, adj) -> list[tuple[str, str, int]]:
    """BFS path as a list of (from_node, edge_id, sign) hops from start to goal."""
    if start == goal:
        return []
    prev: dict[str, tuple[str, str, int]] = {}
    seen = {start}
    q = deque([start])
    while q:
        cur = q.popleft()
        for nb, eid, sign in adj.get(cur, []):
            if nb not in seen:
                seen.add(nb)
                prev[nb] = (cur, eid, sign)
                if nb == goal:
                    q.clear()
                    break
                q.append(nb)
    if goal not in prev:
        return []
    hops: list[tuple[str, str, int]] = []
    node = goal
    while node != start:
        frm, eid, sign = prev[node]
        hops.append((frm, eid, sign))
        node = frm
    hops.reverse()
    return hops


def _accumulate_along_path(path, adj, amount, result) -> None:
    for _frm, eid, sign in path:
        result.edge_flow[eid] += sign * amount


def apply_to_document(document) -> FlowResult:
    """Solve and write ``flow`` / ``flow_direction`` back onto each edge."""
    result = solve(document)
    for e in document.edges():
        f = result.edge_flow.get(e.edge_id, 0.0)
        e.flow = abs(f)
        e.flow_direction = result.direction(e.edge_id)
    return result