AutoPumping 0.1.0__py3-none-any.whl

auto_pumping/executor.py

@@ -0,0 +1,178 @@
+from threading import Event, Thread
+from time import sleep
+
+
+class ConfirmationRequest:
+    """A pending request for the operator to approve a single plan step.
+
+    Created by the executor thread and resolved from the UI event loop, so the
+    handshake goes through a :class:`threading.Event`.
+    """
+
+    def __init__(self, step):
+        self.step = step
+        self._event = Event()
+        self._approved = False
+
+    @property
+    def resolved(self):
+        return self._event.is_set()
+
+    def resolve(self, approved):
+        self._approved = bool(approved)
+        self._event.set()
+
+    def wait_for_response(self, should_continue=None):
+        """Block until the request is resolved, returning whether it was approved.
+
+        Polls rather than blocking indefinitely so that a plan cancelled while
+        the dialog is open does not leave this thread stuck forever.
+        """
+        while not self._event.wait(timeout=0.1):
+            if should_continue is not None and not should_continue():
+                return False
+        return self._approved
+
+
+class PlanExecutor:
+    def __init__(self, pumping_system):
+        self._pumping_system = pumping_system
+        self._current_step = None
+        self._plan = None
+        self._thread = None
+        self._pending_confirmation = None
+        self._confirmed_step = None
+
+    @property
+    def plan(self):
+        return ([self._current_step] if self._current_step is not None else []) + (
+            self._plan if self._plan is not None else []
+        )
+
+    @property
+    def pending_confirmation(self):
+        return self._pending_confirmation
+
+    def execute_plan(self, plan):
+        self.cancel()
+        self._plan = plan
+        self._thread = Thread(target=self._plan_executor)
+        self._thread.start()
+
+    def cancel(self):
+        self._plan = None
+        self._current_step = None
+        self._confirmed_step = None
+        # Release a thread parked on a confirmation dialog before joining it,
+        # otherwise cancelling while the dialog is open deadlocks.
+        pending = self._pending_confirmation
+        if pending is not None:
+            pending.resolve(False)
+        if self._thread is not None:
+            self._thread.join()
+            self._thread = None
+        self._pending_confirmation = None
+
+    def _requires_confirmation(self, step):
+        confirm = self._pumping_system._config.confirm
+        if step.wait is not None:
+            return False
+        if step.valve is not None:
+            return self._confirm_flag(
+                confirm.valve, step.valve, step.state, confirm.default
+            )
+        if step.pump is not None:
+            return self._confirm_flag(
+                confirm.pump, step.pump, step.state, confirm.default
+            )
+        if step.stage is not None:
+            # `confirm.stage` lists the positions that need confirmation; any
+            # other position falls back to `confirm.default`.
+            if confirm.stage is not None and step.stage in confirm.stage:
+                return True
+            return confirm.default
+        return False
+
+    @staticmethod
+    def _confirm_flag(mapping, key, state, default):
+        """Resolve a `Confirm` entry for a single target.
+
+        A boolean entry applies to every action on that target, a state entry
+        only applies when the action moves the target to that state, and an
+        absent entry falls back to `confirm.default`.
+        """
+        if mapping is None or key not in mapping:
+            return default
+        value = mapping[key]
+        if isinstance(value, bool):
+            return value
+        return value == state
+
+    def _request_confirmation(self, step):
+        """Ask the operator to approve `step`, blocking until they answer.
+
+        A step may be revisited by the loop below while it waits for a safety
+        condition, so the answer is remembered against that step rather than
+        being asked again on every pass.
+        """
+        if self._confirmed_step is step:
+            return True
+        if not self._requires_confirmation(step):
+            self._confirmed_step = step
+            return True
+        request = ConfirmationRequest(step)
+        self._pending_confirmation = request
+        try:
+            approved = request.wait_for_response(lambda: self._plan is not None)
+        finally:
+            self._pending_confirmation = None
+        if approved:
+            self._confirmed_step = step
+        return approved
+
+    def _wait_satisfied(self, step):
+        return all(
+            ineq.evaluate(self._pumping_system.get_gauge_pressure(gauge))
+            for gauge, ineq in step.wait.items()
+        )
+
+    def _plan_executor(self):
+        while self._plan is not None and len(self._plan) > 0:
+            if self._current_step is None:
+                self._current_step = self._plan.pop(0)
+            step = self._current_step
+            if step.wait is not None:
+                if self._wait_satisfied(step):
+                    self._current_step = None
+                    continue
+            elif step.stage is not None and self._pumping_system._safety_machine.can_move_stage(step.stage):
+                if self._request_confirmation(step):
+                    self._pumping_system._move_stage(step.stage)
+                    self._current_step = None
+                    continue
+                else:
+                    self._plan = None
+                    self._current_step = None
+                    self._confirmed_step = None
+                    return
+            elif step.valve is not None and self._pumping_system._safety_machine.can_actuate_valve(step.valve, step.state):
+                if self._request_confirmation(step):
+                    self._pumping_system._actuate_valve(step.valve, step.state)
+                    self._current_step = None
+                    continue
+                else:
+                    self._plan = None
+                    self._current_step = None
+                    self._confirmed_step = None
+                    return
+            elif step.pump is not None and self._pumping_system._safety_machine.can_set_pump(step.pump, step.state):
+                if self._request_confirmation(step):
+                    self._pumping_system._set_pump(step.pump, step.state)
+                    self._current_step = None
+                    continue
+                else:
+                    self._plan = None
+                    self._current_step = None
+                    self._confirmed_step = None
+                    return
+            sleep(1)

auto_pumping/graph.py

@@ -0,0 +1,358 @@
+import re
+
+from networkx import DiGraph, dijkstra_path, dijkstra_path_length, relabel_nodes
+from yaml import Node
+from .config import ValveState, PumpState, Inequality, StrictModel
+from pydantic import model_validator
+from uuid import uuid4, UUID
+from math import inf
+from pyvis.network import Network
+from http.server import BaseHTTPRequestHandler, HTTPServer
+import webbrowser
+
+
+class NodeModel(StrictModel):
+    name: str | None = None
+    valves: dict[str, ValveState] = {}
+    pumps: dict[str, PumpState] = {}
+    stage: str | None = None
+    gauges: dict[str, tuple[float, float]] = {}
+
+
+class EdgeModel(StrictModel):
+    time: float
+    wait: dict[str, Inequality] = {}
+
+
+class PlanStep(StrictModel):
+    wait: dict[str, Inequality] | None = None
+    stage: str | None = None
+    valve: str | None = None
+    pump: str | None = None
+    state: ValveState | PumpState | None = None
+
+    @model_validator(mode="after")
+    def validate_wait(self):
+        if self.wait is not None:
+            if self.stage is not None:
+                raise ValueError("Cannot specify both wait and stage")
+            if self.valve is not None:
+                raise ValueError("Cannot specify both wait and valve")
+            if self.pump is not None:
+                raise ValueError("Cannot specify both wait and pump")
+            if self.state is not None:
+                raise ValueError("Cannot specify both wait and state")
+        return self
+
+    @model_validator(mode="after")
+    def validate_state(self):
+        if self.state is not None:
+            if self.valve is None and self.pump is None:
+                raise ValueError("State specified without valve or pump")
+            if self.valve is not None and self.pump is not None:
+                raise ValueError("Cannot specify both valve and pump with state")
+            if self.valve is not None and not isinstance(self.state, ValveState):
+                raise ValueError("State must be a ValveState when valve is specified")
+            if self.pump is not None and not isinstance(self.state, PumpState):
+                raise ValueError("State must be a PumpState when pump is specified")
+        return self
+
+
+class PumpingGraph:
+    def __init__(self, config):
+        self._gauge_ranges = self._discretize_gauge_ranges(
+            config.gauges, config.transitions
+        )
+        self._graph = DiGraph()
+        self._add_initial_states(config.initial_states, config.named_states)
+        self._add_reachable_states(config.transitions, config.named_states)
+
+    def _get_gauge_range(self, gauge, pressure):
+        for gauge_range in self._gauge_ranges[gauge]:
+            if gauge_range[0] <= pressure < gauge_range[1]:
+                return gauge_range
+        raise ValueError(f"Pressure {pressure} is out of range for gauge {gauge}")
+
+    def _add_initial_states(self, initial_states, named_states):
+        for state in initial_states:
+            node_data = NodeModel(
+                name=self._get_state_name(state, named_states),
+                valves=state.valves,
+                pumps=state.pumps,
+                stage=state.stage,
+                gauges={
+                    gauge: self._get_gauge_range(gauge, pressure)
+                    for gauge, pressure in state.gauges.items()
+                },
+            )
+            self._graph.add_node(
+                uuid4(),
+                **node_data.model_dump(),
+            )
+
+    def _add_reachable_states(self, transitions, named_states):
+        for state_id in list(self._graph.nodes):
+            self._add_reachable_states_from(state_id, transitions, named_states)
+
+    def _get_state_data(self, state_id):
+        return NodeModel(**self._graph.nodes[state_id])
+
+    def _add_state(self, state_data):
+        assert self._find_state(state_data) is None
+        state_id = uuid4()
+        self._graph.add_node(state_id, **state_data.model_dump())
+        return state_id
+
+    def _add_reachable_states_from(self, state_id, transitions, named_states):
+        state_data = self._get_state_data(state_id)
+        for transition in transitions:
+            if not self._is_valid_transition(state_data, transition):
+                continue
+            reachable_states = self._build_reachable_states(
+                state_data, transition, named_states
+            )
+            for new_state_data, time in reachable_states:
+                new_state_id = self._find_state(new_state_data)
+                if new_state_id is None:
+                    new_state_id = self._add_state(new_state_data)
+                    self._add_reachable_states_from(
+                        new_state_id, transitions, named_states
+                    )
+                self._add_edge(state_id, new_state_id, transition, time)
+
+    def _add_edge(self, from_state_id, to_state_id, transition, time):
+        if not self._graph.has_edge(from_state_id, to_state_id):
+            edge_data = EdgeModel(
+                wait=transition.wait,
+                time=time,
+            )
+            self._graph.add_edge(
+                from_state_id,
+                to_state_id,
+                **edge_data.model_dump(),
+            )
+
+    @staticmethod
+    def _is_valid_transition(state_data, transition):
+        if transition.stage and state_data.stage != transition.from_state:
+            return False
+        if (
+            transition.valve is not None
+            and state_data.valves.get(transition.valve) != transition.from_state
+        ):
+            return False
+        if (
+            transition.pump is not None
+            and state_data.pumps.get(transition.pump) != transition.from_state
+        ):
+            return False
+        if transition.condition is None:
+            return True
+        for gauge, ineq in transition.condition.gauge.items():
+            pressure_range = state_data.gauges[gauge]
+            pressure = sum(pressure_range) / 2
+            if not ineq.evaluate(pressure):
+                return False
+        for valve, valve_condition in transition.condition.valve.items():
+            valve_state = valve_condition.state
+            if state_data.valves.get(valve) != valve_state:
+                return False
+        for pump, pump_state in transition.condition.pump.items():
+            if state_data.pumps.get(pump) != pump_state.state:
+                return False
+        if (
+            transition.condition.stage is not None
+            and state_data.stage != transition.condition.stage
+        ):
+            return False
+        return True
+
+    def _find_state(self, state_data):
+        for node_id in self._graph.nodes:
+            node_data = self._get_state_data(node_id)
+            if (
+                node_data.valves == state_data.valves
+                and node_data.pumps == state_data.pumps
+                and node_data.stage == state_data.stage
+                and node_data.gauges == state_data.gauges
+            ):
+                return node_id
+        return None
+
+    def _find_states(self, name=None, valves=None, pumps=None, stage=None):
+        for node_id in self._graph.nodes:
+            node_data = self._get_state_data(node_id)
+            if name is not None and node_data.name != name:
+                continue
+            if valves is not None and any(
+                state != node_data.valves.get(valve) for valve, state in valves.items()
+            ):
+                continue
+            if pumps is not None and any(
+                state != node_data.pumps.get(pump) for pump, state in pumps.items()
+            ):
+                continue
+            if stage is not None and node_data.stage != stage:
+                continue
+            yield node_id
+
+    def _build_reachable_states(self, state_data, transition, named_states):
+        new_state_data = state_data.model_copy(deep=True)
+        if transition.valve is not None:
+            new_state_data.valves[transition.valve] = transition.to_state
+        if transition.pump is not None:
+            new_state_data.pumps[transition.pump] = transition.to_state
+        if transition.stage:
+            new_state_data.stage = transition.to_state
+        new_state_data.name = self._get_state_name(new_state_data, named_states)
+        yield new_state_data, 1
+        for pressure_change in transition.pressure:
+            new_state_data = new_state_data.model_copy(deep=True)
+            new_state_data.gauges[pressure_change.gauge] = self._get_gauge_range(
+                pressure_change.gauge, pressure_change.pressure
+            )
+            new_state_data.name = self._get_state_name(new_state_data, named_states)
+            yield new_state_data, 60 * pressure_change.time
+
+    @staticmethod
+    def _get_state_name(state, named_states):
+        for name, named_state in named_states.items():
+            is_named = True
+            for valve, valve_state in named_state.valves.items():
+                if state.valves.get(valve) != valve_state:
+                    is_named = False
+                    break
+            for pump, pump_state in named_state.pumps.items():
+                if state.pumps.get(pump) != pump_state:
+                    is_named = False
+                    break
+            if named_state.stage is not None and state.stage != named_state.stage:
+                is_named = False
+            if is_named:
+                return name
+        return None
+
+    @staticmethod
+    def _discretize_gauge_ranges(gauges, transitions):
+        gauge_range_edges = {gauge: [0, inf] for gauge in gauges}
+        for transition in transitions:
+            if transition.condition is None:
+                continue
+            if transition.condition.gauge is None:
+                continue
+            for gauge, ineq in transition.condition.gauge.items():
+                if ineq.LE is not None and ineq.LE not in gauge_range_edges[gauge]:
+                    gauge_range_edges[gauge].append(ineq.LE)
+                if ineq.GE is not None and ineq.GE not in gauge_range_edges[gauge]:
+                    gauge_range_edges[gauge].append(ineq.GE)
+        for edges in gauge_range_edges.values():
+            edges.sort()
+        return {
+            gauge: [(edges[i], edges[i + 1]) for i in range(len(edges) - 1)]
+            for gauge, edges in gauge_range_edges.items()
+        }
+
+    @staticmethod
+    def _state_diff(start_state_data, end_state_data):
+        if start_state_data.stage != end_state_data.stage:
+            return PlanStep(stage=end_state_data.stage)
+        for valve, state in end_state_data.valves.items():
+            if start_state_data.valves.get(valve) != state:
+                return PlanStep(valve=valve, state=state)
+        for pump, state in end_state_data.pumps.items():
+            if start_state_data.pumps.get(pump) != state:
+                return PlanStep(pump=pump, state=state)
+        raise ValueError("No difference between states")
+
+    def make_plan(self, start_state, **goal_state_data):
+        if not isinstance(start_state, NodeModel):
+            start_state = NodeModel(**start_state)
+        if not isinstance(start_state, UUID):
+            start_state = self._find_state(start_state)
+        assert start_state is not None, f"Cannot find state {start_state}"
+        goal_lengths = sorted(
+            [
+                (
+                    dijkstra_path_length(self._graph, start_state, goal_state, "time"),
+                    goal_state,
+                )
+                for goal_state in self._find_states(**goal_state_data)
+            ],
+            key=lambda x: x[0],
+        )
+        assert len(goal_lengths), "No goal state found matching criteria"
+        path = dijkstra_path(self._graph, start_state, goal_lengths[0][1], "time")
+        for i in range(len(path) - 1):
+            edge_data = self._graph.get_edge_data(path[i], path[i + 1])
+            start_node_data = self._get_state_data(path[i])
+            end_node_data = self._get_state_data(path[i + 1])
+            if edge_data["wait"]:
+                yield PlanStep(wait=edge_data["wait"])
+            yield self._state_diff(start_node_data, end_node_data)
+
+    def draw(self, output_file=None, port=8000):
+        net = Network(height="100vh", width="100%", directed=True)
+        graph = relabel_nodes(self._graph, str, copy=True)
+        for node in graph.nodes:
+            node_data = self._get_state_data(UUID(node))
+            title = ""
+            title += f"Stage: {node_data.stage}\n"
+            title += "Valves:\n"
+            for valve, state in node_data.valves.items():
+                title += f"  {valve}: {state.value}\n"
+            title += "Pumps:\n"
+            for pump, state in node_data.pumps.items():
+                title += f"  {pump}: {state.value}\n"
+            title += "Gauges:\n"
+            for gauge, pressure_range in node_data.gauges.items():
+                title += f"  {gauge}: {pressure_range[0]} - {pressure_range[1]}\n"
+            graph.nodes[node]["label"] = node_data.name if node_data.name else " "
+            graph.nodes[node]["title"] = title.strip()
+        for edge in graph.edges:
+            edge_data = EdgeModel(**graph.edges[edge])
+            title = f"Time: {edge_data.time} s\n"
+            if edge_data.wait:
+                title += "Wait:\n"
+                for gauge, ineq in edge_data.wait.items():
+                    title += f"  {gauge}: {ineq}\n"
+            graph.edges[edge]["title"] = title.strip()
+            edge_diff = self._state_diff(
+                self._get_state_data(UUID(edge[0])), self._get_state_data(UUID(edge[1]))
+            )
+            if edge_diff.stage is not None:
+                graph.edges[edge]["label"] = f"to {edge_diff.stage}"
+            if edge_diff.valve is not None:
+                graph.edges[edge][
+                    "label"
+                ] = f"{edge_diff.valve} {edge_diff.state.value}"
+            if edge_diff.pump is not None:
+                graph.edges[edge]["label"] = f"{edge_diff.pump} {edge_diff.state.value}"
+        net.from_nx(graph)
+        if output_file:
+            net.write_html(output_file)
+            return
+        page = net.generate_html().encode("utf-8")
+
+        class InMemoryGraphHandler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                self.send_response(200)
+                self.send_header("Content-type", "text/html")
+                self.send_header("Content-length", str(len(page)))
+                self.end_headers()
+                self.wfile.write(page)
+
+            def log_message(self, format, *args):
+                return  # Suppress logging
+
+        HTTPServer.allow_reuse_address = True
+        server = HTTPServer(("127.0.0.1", port), InMemoryGraphHandler)
+        print(f"Server at http://127.0.0.1:{port}")
+
+        # Open the browser automatically
+        webbrowser.open(f"http://127.0.0.1:{port}")
+
+        try:
+            server.serve_forever()
+        except KeyboardInterrupt:
+            print("\nShutting down server.")
+            server.server_close()

auto_pumping/safety.py

@@ -0,0 +1,83 @@
+from time import time
+
+
+class SafetyMachine:
+    def __init__(self, config, pumping_system):
+        self._config = config
+        self._pumping_system = pumping_system
+        self._valve_times = {valve: time() for valve in config.valves}
+        self._pump_times = {pump: time() for pump in config.pumps}
+
+    def record_valve_change(self, valve):
+        self._valve_times[valve] = time()
+
+    def record_pump_change(self, pump):
+        self._pump_times[pump] = time()
+
+    def can_actuate_valve(self, valve, state):
+        current_state = self._pumping_system.get_valve_state(valve)
+        transition = self._find_transition(
+            valve=valve, from_state=current_state, to_state=state
+        )
+        return self._check_conditions(transition.condition)
+
+    def can_set_pump(self, pump, state):
+        current_state = self._pumping_system.get_pump_state(pump)
+        transition = self._find_transition(
+            pump=pump, from_state=current_state, to_state=state
+        )
+        return self._check_conditions(transition.condition)
+
+    def can_move_stage(self, position):
+        current_position = self._pumping_system.get_stage_position()
+        transition = self._find_transition(
+            stage=True, from_state=current_position, to_state=position
+        )
+        return self._check_conditions(transition.condition)
+
+    def _find_transition(self, **kwargs):
+        for transition in self._config.transitions:
+            if all(getattr(transition, key) == value for key, value in kwargs.items()):
+                return transition
+        raise ValueError(f"No transition found for {kwargs}")
+
+    def _check_conditions(self, conditions):
+        if conditions is None:
+            return True
+        for gauge, ineq in conditions.gauge.items():
+            if not ineq.evaluate(self._pumping_system.get_gauge_pressure(gauge)):
+                return False
+        if (
+            conditions.stage is not None
+            and conditions.stage != self._pumping_system.get_stage_position()
+        ):
+            return False
+        for valve, condition in conditions.valve.items():
+            if self._pumping_system.get_valve_state(valve) != condition.state:
+                return False
+            if (
+                condition.time is not None
+                and time() - self._valve_times[valve] < condition.time
+            ):
+                return False
+        for pump, condition in conditions.pump.items():
+            if condition.state is not None and self._pumping_system.get_pump_state(pump) != condition.state:
+                return False
+            if (
+                condition.time is not None
+                and time() - self._pump_times[pump] < condition.time
+            ):
+                return False
+            if condition.current is not None and not condition.current.evaluate(
+                self._pumping_system.get_pump_current(pump)
+            ):
+                return False
+            if condition.voltage is not None and not condition.voltage.evaluate(
+                self._pumping_system.get_pump_voltage(pump)
+            ):
+                return False
+            if condition.flow is not None and not condition.flow.evaluate(self._pumping_system.get_pump_flow(pump)):
+                return False
+            if condition.speed is not None and not condition.speed.evaluate(self._pumping_system.get_pump_speed(pump)):
+                return False
+        return True