diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index ecba9eb5..1b79f7e4 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -17,6 +17,14 @@ jobs: - name: Checkout uses: actions/checkout@v6 + - name: Download CPLEX installer from Azure storage account + uses: armanrahman22/azblob-download-action@v0.0.4 + with: + connection-string: ${{ secrets.BLDDEPS_STORAGE_ACCOUNT_CONN_STR }} + container-name: "solver-blobs" + blob-name: "cplex_22.1.1.0-1_amd64.deb" + download-path: "." + - name: Install uv uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 with: @@ -24,13 +32,18 @@ jobs: - name: Install Dependencies run: | - uv sync + uv sync --group dev + sudo apt install -f ./cplex_22.1.1.0-1_amd64.deb - name: Coverage (source branch) run: | echo "[run]" > .coveragecfg echo "data_file=.coverage-source" >> .coveragecfg - uv run pytest -n auto --cov=my_app --cov-config=.coveragecfg . + uv run pytest -n auto --cov=echo --cov-config=.coveragecfg . + env: + OPTIMISER_ENGINE: "cplex" + OPTIMISER_ENGINE_EXECUTABLE: "/opt/cplex/cplex" + PYTHONPATH: ./ - name: Extract Coverage id: extract @@ -57,6 +70,14 @@ jobs: with: ref: ${{ github.base_ref }} + - name: Download CPLEX installer from Azure storage account + uses: armanrahman22/azblob-download-action@v0.0.4 + with: + connection-string: ${{ secrets.BLDDEPS_STORAGE_ACCOUNT_CONN_STR }} + container-name: "solver-blobs" + blob-name: "cplex_22.1.1.0-1_amd64.deb" + download-path: "." + - name: Install uv uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 with: @@ -64,13 +85,18 @@ jobs: - name: Install Dependencies run: | - uv sync + uv sync --group dev + sudo apt install -f ./cplex_22.1.1.0-1_amd64.deb - name: Coverage (target branch) run: | echo "[run]" > .coveragecfg echo "data_file=.coverage-target" >> .coveragecfg - uv run pytest -n auto --cov=my_app --cov-config=.coveragecfg . + uv run pytest -n auto --cov=echo --cov-config=.coveragecfg . + env: + OPTIMISER_ENGINE: "cplex" + OPTIMISER_ENGINE_EXECUTABLE: "/opt/cplex/cplex" + PYTHONPATH: ./ - name: Extract Coverage id: extract diff --git a/examples/attribute_tracking.py b/examples/attribute_tracking.py new file mode 100644 index 00000000..7d583397 --- /dev/null +++ b/examples/attribute_tracking.py @@ -0,0 +1,201 @@ +import numpy as np +import pandas as pd +from typing import Any +import networkx as nx + +from echo.configuration import Units +from echo.models.agnostic import FlexPort, Sink, TellegenNode +from echo.models.base import Node, OptimisationGraph +from echo.models.scenario import ScenarioSettings, engine_settings_from_environment +from echo.models.thermal import ThermalStorage +from echo.objectives.base import ObjectiveSet +from echo.objectives.power import PeakPositivePower +from echo.objectives.tariff import ThroughputCost +from echo.optimiser import optimise +from echo.utils import TimeSeriesData, expand_as_dict +from echo import visualization as viz + +pd.options.plotting.backend = "plotly" + + +""" A simple thermal network. + A heating load, thermal storage and heating mains are connected via a connection point (TellegenNode). +""" + +# ---------------------------------------------------------------------------------------------------------------------- +# 1. Define constants +# ---------------------------------------------------------------------------------------------------------------------- + +NUMBER_INTERVALS = 48 +INTERVAL_DURATION = 30 +NUM_EXPANSION_PERIODS = 1 + + +def default_surface_area_of_cylinder(volume: float, include_bottom: bool = True): + """Given volume of the cylinder in cubic meters, calculate surface area. Assuming height to diameter ration H/D=3. + + If include_bottom is False, do not include bottom surface. + """ + radius = np.cbrt(volume / (np.pi * 6)) + height = 6 * radius + if include_bottom: + return round(2 * np.pi * radius * height + 2 * np.pi * radius**2, 3) + else: + return round(2 * np.pi * radius * height + np.pi * radius**2, 3) + + +# Thermal transmittance of storage insulation in W/sqm*C (from 0.5 - 11 is reasonable range) +INSULATION_TRANSMITTANCE = 5 +# mass of thermal storage in kg +MASS = 500 +# Specific heat capacity of storage medium (here is water) in J/kg*C +SPECIFIC_HEAT_CAPACITY_WATER = 4184 +# Total surface area of thermal storage from volume (for water 1 kg=1 litre) +SURFACE_AREA = default_surface_area_of_cylinder(MASS * 1e-3) + +# ---------------------------------------------------------------------------------------------------------------------- +# 2. Define thermal demand profile and ambient temperature profile +# ---------------------------------------------------------------------------------------------------------------------- + +q_max_joules = SPECIFIC_HEAT_CAPACITY_WATER * MASS * 70 # Max energy storage capacity in joules +q_max_kwh = q_max_joules / 3600000 +th_load = [0.1] * 14 + [0.4] * 4 + [0.05] * 16 + [0.4] * 6 + [0.2] * 8 +th_load = list((np.array(th_load) * q_max_kwh).round()) + +th_demand_data = TimeSeriesData( + value=th_load, num_time_intervals=NUMBER_INTERVALS, num_expansion_intervals=NUM_EXPANSION_PERIODS +) + +th_demand_dict = expand_as_dict(th_demand_data) + +amb_temp_data = TimeSeriesData( + value=25, num_time_intervals=NUMBER_INTERVALS, num_expansion_intervals=NUM_EXPANSION_PERIODS +) +ambient_temp_dict = expand_as_dict(amb_temp_data) + +# ---------------------------------------------------------------------------------------------------------------------- +# 3. Define nodes +# ---------------------------------------------------------------------------------------------------------------------- + +thermal_demand = Node(node_name="thermal_load", ports={"demand_kwt": Sink(units=Units.KWT)}) + +thermal_demand.ports["demand_kwt"].add_sink_profile(th_demand_dict) + +thermal_mains = Node(node_name="thermal_supply", ports={"supply_kwt": FlexPort(units=Units.KWT)}) + +storage = ThermalStorage( + node_name="thermal_storage", + max_temp=80, + min_temp=10, + ambient_temp=ambient_temp_dict, + storage_mass=MASS, + specific_heat=SPECIFIC_HEAT_CAPACITY_WATER, + ins_transmittance=INSULATION_TRANSMITTANCE, + surface_area=SURFACE_AREA, + separate_in_out_ports=False, +) + +cp = TellegenNode( + node_name="conn_point", + ports={ + "to_supply_kwt": FlexPort(units=Units.KWT), + "to_storage_kwt": FlexPort(units=Units.KWT), + "to_demand_kwt": FlexPort(units=Units.KWT), + }, +) + +# ---------------------------------------------------------------------------------------------------------------------- +# 4. Build the optimisation graph +# ---------------------------------------------------------------------------------------------------------------------- +# +# thermal_mains---cp---thermal_demand +# | +# storage +# +system = OptimisationGraph() +system.add_node_obj([storage, thermal_demand, thermal_mains, cp]) +system.connect_ports_and_create_edge(cp.ports["to_supply_kwt"], thermal_mains.ports["supply_kwt"]) +system.connect_ports_and_create_edge(cp.ports["to_storage_kwt"], storage.ports["input_output"]) +system.connect_ports_and_create_edge(cp.ports["to_demand_kwt"], thermal_demand.ports["demand_kwt"]) + +objective_set = ObjectiveSet( + objective_list=[ + ThroughputCost(component=storage.ports["input_output"], rate=0.01), + PeakPositivePower(component=cp.ports["to_supply_kwt"]), + ] +) + + +optimise_results = optimise( + scenario_settings=ScenarioSettings( + interval_duration=INTERVAL_DURATION, + number_of_intervals=NUMBER_INTERVALS, + number_of_expansion_intervals=NUM_EXPANSION_PERIODS, + ), + engine_settings=engine_settings_from_environment(), + graph=system, + objective_set=objective_set, +) + +optimise_results.model.pprint() +# optimise_results.model.display() + + +graph = system.convert_to_nx() +positions = nx.spring_layout(graph) + +tracker = optimise_results.model_attribute_tracker +print(tracker) + + +def to_html_string(attrs: dict[str, list[Any]]): + attr_text = "" + for attr_class in sorted(attrs.keys()): + attr_text += "
" + attr_descriptions = [f" {attr_class} '{attr_name}'
" for attr_name in sorted(attrs[attr_class])] + attr_text += "".join(attr_descriptions) + + return attr_text + + +def node_hover_text(node: Any, *_) -> str: + node_object = optimise_results.graph.node_obj[node] + node_type = type(node_object).__name__ + + tracker = optimise_results.model_attribute_tracker + if tracker is not None: + attrs = tracker.filtered_pairwise_attributes(include=lambda _, checkpoint: checkpoint.startswith(node)) + title = f"{node_type} '{node}'
" + return title + to_html_string(attrs) + + return "" + + +def edge_hover_text(edge: Any, *_) -> str: + edge_reversed = edge[::-1] # reverse the edge tuple (not sure why edges are defined in the opposite sense) + edge_object = optimise_results.graph.edge_obj[edge_reversed] + + tracker = optimise_results.model_attribute_tracker + if tracker is not None: + # Get the attributes that were added when `edge` was added to the pyomo model + attrs = tracker.filtered_pairwise_attributes(include=lambda _, checkpoint: checkpoint == edge_object.edge_name) + title = f"{edge}
" + return title + to_html_string(attrs) + return "" + +def node_size(*_) -> int: + return 40 + +# Use the new visualization function to plot the network/graph +# This function takes callbacks for setting the node (colors, text/labels, hover text etc). +fig = viz.PlotlyGraph.plot( + graph=graph, + positions=positions, + show_node_names=True, + node_color=viz.echo_node_color, + node_size=node_size, + node_text=viz.echo_node_text, + node_hover_text=node_hover_text, + edge_hover_text=edge_hover_text, +) +fig.show() diff --git a/src/echo/models/base.py b/src/echo/models/base.py index cf5fa295..0f15b0cb 100644 --- a/src/echo/models/base.py +++ b/src/echo/models/base.py @@ -66,7 +66,7 @@ class Port(BaseModel): export_constraint_value: ConstraintValueType | None = None active_periods: dict[tuple[int, int], Any] | None = None slack: bool = False - objective: float | en.numeric_expr.NumericExpression = 0 # this will eventually be a pyomo expression + objective: en.numeric_expr.NumericExpression | float = 0 # this will eventually be a pyomo expression # Validators for import/export constraint values import_con_sign = validator("import_constraint_value", allow_reuse=True)(import_cons_check) @@ -602,7 +602,7 @@ class Node(BaseModel): node_name: str = "" uid: str = Field(default_factory=shortuuid.uuid) ports: dict[str, Port] = {} - objective: float | en.numeric_expr.NumericExpression = 0 # For adding any node objectives + objective: en.numeric_expr.NumericExpression | float = 0 # For adding any node objectives @property def inflow(self) -> str: @@ -1003,7 +1003,7 @@ class Path(BaseModel): path_name: str | None = None units = Units.KW regularise: bool = False - objective: float | en.numeric_expr.NumericExpression = 0 + objective: en.numeric_expr.NumericExpression | float = 0 flow_value: str = "" contingency_neg: str | None diff --git a/src/echo/optimiser.py b/src/echo/optimiser.py index ce73169e..5c4c9f03 100644 --- a/src/echo/optimiser.py +++ b/src/echo/optimiser.py @@ -12,6 +12,7 @@ from echo.models.base import Node, OptimisationGraph, Port from echo.models.scenario import EchoConcreteModel, EngineSettings, ScenarioSettings from echo.objectives.base import Objective, ObjectiveSet +from echo.tracker import AttributeTracker # The default set of termination conditions that pyomo can return and echo will report as a success DEFAULT_ACCEPTABLE_TERMINATION_CONDITIONS: Collection[TerminationCondition] = set( @@ -45,12 +46,13 @@ class OptimisationResult: """ scenario_settings: ScenarioSettings - objective: en.numeric_expr.NumericExpression + objective: en.numeric_expr.NumericExpression | None model: EchoConcreteModel objective_set: ObjectiveSet | None graph: OptimisationGraph opt_status: SolverStatus termination_condition: TerminationCondition + model_attribute_tracker: AttributeTracker | None = None def df(self) -> pd.DataFrame: """ @@ -184,15 +186,26 @@ def build_model_and_objective( big_m: int, profile: pd.DataFrame | None, objective_set: ObjectiveSet | None, -) -> tuple[EchoConcreteModel, en.numeric_expr.NumericExpression]: +) -> tuple[EchoConcreteModel, en.numeric_expr.NumericExpression | None, AttributeTracker]: """Builds an EchoConcreteModel for a particular Echo Scenario definition and a related objective to optimise against the model""" - model = _build_model( - graph=graph, scenario_settings=scenario_settings, small_m=small_m, big_m=big_m, profile=profile + model = EchoConcreteModel() + tracker = AttributeTracker(model) + + _build_model( + model=model, + graph=graph, + scenario_settings=scenario_settings, + small_m=small_m, + big_m=big_m, + profile=profile, + tracker=tracker, + ) + objective = _build_objective( + model=model, graph=graph, objective_set=objective_set, profile=profile, tracker=tracker ) - objective = _build_objective(model=model, graph=graph, objective_set=objective_set, profile=profile) - return model, objective + return model, objective, tracker def _build_model( @@ -201,10 +214,16 @@ def _build_model( small_m: float, big_m: int, profile: pd.DataFrame | None, + tracker: AttributeTracker, + model: EchoConcreteModel | None = None, ) -> EchoConcreteModel: - model = EchoConcreteModel() + if model is None: + model = EchoConcreteModel() model.small_m = en.Param(initialize=small_m) model.big_m = en.Param(initialize=big_m) + + tracker.mark("model-init") + model.scenario_settings = scenario_settings # We use RangeSet to create an index for each of the time @@ -223,25 +242,29 @@ def _build_model( discount_rates[ep] = 1 / ((1 + scenario_settings.discount_rate) ** ep) model.discount_rates = en.Param(model.Expansion, initialize=discount_rates) + tracker.mark("model-setup") + # Initialise node variables/params and add node constraints for node_obj in graph.node_obj.values(): node_obj.verify_node() node_obj.add_node_to_model(model, profile) + node_obj.apply_node_constraints(model) + tracker.mark(f"{node_obj.node_name}:node-constraints") # Initialise edge variables/params and add edge constraints for edge_obj in graph.edge_obj.values(): edge_obj.verify_edge() edge_obj.add_edge_to_model(model) + tracker.mark(edge_obj.edge_name) # Initialise paths for path in graph.paths.values(): path.add_path_to_model(model) + tracker.mark(path.path_name) - # Apply constraints - for obj in graph.node_obj.values(): - obj.apply_node_constraints(model) if graph.paths: graph.apply_path_constraints(model) + tracker.mark("model-adding-path-constraints") return model @@ -251,13 +274,15 @@ def _build_objective( graph: OptimisationGraph, profile: pd.DataFrame | None, objective_set: ObjectiveSet | None, -) -> float | en.numeric_expr.NumericExpression: - objective: float | en.numeric_expr.NumericExpression = 0 + tracker: AttributeTracker, + objective: en.numeric_expr.NumericExpression | float = 0, +) -> en.numeric_expr.NumericExpression | float | None: # Add objectives defined in the objective set if objective_set is not None: objective_set.add_objectives_to_model(model, profile) objective += objective_set.get_objective_total(model) + tracker.mark("model-adding-objective-set") # Add any other costs that are defined on graph nodes/ports/paths for node_obj in graph.node_obj.values(): @@ -266,11 +291,17 @@ def _build_objective( for port_obj in node_obj.ports.values(): port_obj.add_objective(model) # populate the .objective attribute for each port objective += port_obj.objective # add the newly populated attribute to our total + tracker.mark(f"{node_obj.node_name}:objectives") for path_obj in graph.paths.values(): path_obj.add_objective(model) objective += path_obj.objective + tracker.mark(path_obj.path_name if path_obj.path_name is not None else "model-adding-path-objectives") + + # Determine if we failed to build an objective + if isinstance(objective, float) and objective == 0: + return None return objective @@ -299,7 +330,7 @@ def optimise( validate_network_graph(graph) - model, objective = build_model_and_objective( + model, objective, tracker = build_model_and_objective( graph=graph, scenario_settings=scenario_settings, small_m=engine_settings.small_m, @@ -308,10 +339,12 @@ def optimise( objective_set=objective_set, ) - def objective_function(model: EchoConcreteModel) -> en.numeric_expr.NumericExpression: - return objective + if objective is not None: + + def cost_function(model: EchoConcreteModel) -> en.numeric_expr.NumericExpression: + return objective - model.total_cost = en.Objective(rule=objective_function, sense=en.minimize) + model.total_cost = en.Objective(rule=cost_function, sense=en.minimize) # Set the path to the solver if engine_settings.engine_executable: @@ -361,4 +394,5 @@ def objective_function(model: EchoConcreteModel) -> en.numeric_expr.NumericExpre graph=graph, opt_status=solver_status, termination_condition=termination_condition, + model_attribute_tracker=tracker, ) diff --git a/src/echo/tracker.py b/src/echo/tracker.py new file mode 100644 index 00000000..e103078d --- /dev/null +++ b/src/echo/tracker.py @@ -0,0 +1,74 @@ +from collections import defaultdict +from collections.abc import Callable +from itertools import pairwise +from typing import Any + + +class AttributeTracker: + """Attribute Tracker + + Each checkpoint records *all* the attributes on the object at the checkpoint. + + Therefore to see which attributes were added between two checkpoints, we need + to calculate the difference, which we can do with the `diff` method. + + Usage: + >>> class Foo(object): + ... pass + ... + >>> foo = Foo() + >>> tracker = AttributeTracker() + >>> tracker.track(foo, "init") + >>> foo.a = 1 + >>> tracker.track(foo, "setting-a") + >>> tracker.diff("setting-a", "init") + >>> attrs = tracker.filtered_pairwise_attributes(include: lambda prev_k, k: k=="checkpoint-of-interest") + >>> tracker.to_html_string(foo, "setting-a") + """ + + def __init__(self, object: object) -> None: + """Store the attributes for a checkpoint. The ordering of the keys (checkpoints) is important""" + self.attributes: dict[str, list[str]] = {} + self.object = object + + def mark(self, checkpoint: str) -> None: + """Create a new checkpoint and associate it with the objects attributes.""" + self.attributes[checkpoint] = dir(self.object) + + def diff(self, checkpoint1: str, checkpoint2: str) -> list[str]: + """Returns the attributes that are present at checkpoint2 but not a checkpoint1.""" + s1 = set(self.attributes[checkpoint1]) + s2 = set(self.attributes[checkpoint2]) + + return list(s2 - s1) + + def __str__(self) -> str: + """Returns string representation of all attributes added between first and last checkpoints.""" + result = "" + for prev_checkpoint, checkpoint in pairwise(self.attributes.keys()): + for attr_name in self.diff(prev_checkpoint, checkpoint): + result += f"- {attr_name} ({type(getattr(self.object, attr_name))})\n" + + return result + + def filtered_pairwise_attributes(self, include: Callable[[str, str], bool]) -> dict[str, list[Any]]: + """ + Pairwise diff (attributes added to object between consecutive checkpoints). + Only including pairs where filter function `include(prev_checkpoint, checkpoint)` returns True. + And groups attributes by type (as class-name string). + + Returns a mapping from attribute class name to a list of attributes (by name) belonging to that class. + """ + + attrs = defaultdict(list) + for prev_checkpoint, checkpoint in pairwise(self.attributes.keys()): + if include(prev_checkpoint, checkpoint): + attr_keys = self.diff(prev_checkpoint, checkpoint) + + for attr_key in attr_keys: + attr = getattr(self.object, attr_key) + attr_classname = type(attr).__name__ + + # Group attributes by class + attrs[attr_classname].append(attr_key) + return attrs diff --git a/src/echo/visualization.py b/src/echo/visualization.py new file mode 100644 index 00000000..f29566df --- /dev/null +++ b/src/echo/visualization.py @@ -0,0 +1,229 @@ +from collections.abc import Callable +from enum import StrEnum, auto +from typing import Any + +import networkx as nx +import plotly.graph_objects as go + +DEFAULT_NODE_NAME = "node" +DEFAULT_NODE_SIZE = 15 +DEFAULT_NODE_COLOR = "#888" +DEFAULT_NODE_HOVER_TEXT = "" +DEFAULT_EDGE_TEXT = "" +DEFAULT_EDGE_HOVER_TEXT = "" + + +class VizNodeType(StrEnum): + Bus = auto() + Line = auto() + Grid = auto() + Transformer = auto() + Load = auto() + Solar = auto() + ElectricalLoad = auto() + EV = auto() + Battery = auto() + Inverter = auto() + ConnectionPoint = auto() + HeatPump = auto() + ParameterisedHeatPump = auto() + Chiller = auto() + ParameterisedChiller = auto() + ThermalDistribution = auto() + ThermalTellegen = auto() + ThermalStorage = auto() + HeatAggregationNode = auto() + HeatingDemand = auto() + CoolingDemand = auto() + MultiCommodityBus = auto() + PartitionedMultiCommodityBus = auto() + GasGrid = auto() + Boiler = auto() + + +COLOR_BY_NODE_TYPE = { + VizNodeType.Bus: "darkslategrey", + VizNodeType.Line: "darkslategrey", + VizNodeType.Grid: "hotpink", + VizNodeType.Transformer: "blueviolet", + VizNodeType.Load: "forestgreen", + VizNodeType.Solar: "orange", + VizNodeType.EV: "orangered", + VizNodeType.Battery: "lightblue", + VizNodeType.Inverter: "lightgreen", + VizNodeType.ConnectionPoint: "mediumblue", + VizNodeType.HeatPump: "tomato", + VizNodeType.ParameterisedHeatPump: "tomato", + VizNodeType.Chiller: "cornflowerblue", + VizNodeType.ParameterisedChiller: "cornflowerblue", + VizNodeType.ThermalDistribution: "goldenrod", + VizNodeType.ThermalTellegen: "chocolate", + VizNodeType.ThermalStorage: "firebrick", + VizNodeType.HeatAggregationNode: "firebrick", + VizNodeType.HeatingDemand: "crimson", + VizNodeType.CoolingDemand: "steelblue", + VizNodeType.MultiCommodityBus: "darkblue", + VizNodeType.PartitionedMultiCommodityBus: "lightseagreen", + VizNodeType.GasGrid: "pink", + VizNodeType.Boiler: "firebrick", +} + +SIZE_BY_NODE_TYPE = { + VizNodeType.Bus: 10, + VizNodeType.Line: 10, + VizNodeType.Grid: 30, + VizNodeType.Transformer: 25, + VizNodeType.Load: 15, + VizNodeType.ElectricalLoad: 15, + VizNodeType.Solar: 15, + VizNodeType.EV: 15, + VizNodeType.Battery: 15, + VizNodeType.Inverter: 15, + VizNodeType.ConnectionPoint: 20, + VizNodeType.HeatPump: 15, + VizNodeType.ParameterisedHeatPump: 15, + VizNodeType.Chiller: 15, + VizNodeType.ParameterisedChiller: 15, + VizNodeType.ThermalDistribution: 15, + VizNodeType.ThermalTellegen: 15, + VizNodeType.ThermalStorage: 15, + VizNodeType.HeatAggregationNode: 15, + VizNodeType.HeatingDemand: 15, + VizNodeType.CoolingDemand: 15, + VizNodeType.MultiCommodityBus: 20, + VizNodeType.PartitionedMultiCommodityBus: 20, + VizNodeType.GasGrid: 30, + VizNodeType.Boiler: 15, +} + +ECHO_NODE_TO_VIZ_NODE = { + "thermal_storage": VizNodeType.ThermalStorage, + "thermal_load": VizNodeType.HeatingDemand, + "thermal_supply": VizNodeType.ThermalDistribution, + "conn_point": VizNodeType.ThermalTellegen, +} + + +def echo_node_color(node: str, *args) -> str: # noqa ANN002 + if node in ECHO_NODE_TO_VIZ_NODE: + return COLOR_BY_NODE_TYPE[ECHO_NODE_TO_VIZ_NODE[node]] + return DEFAULT_NODE_COLOR + + +def echo_node_size(node: str, *args) -> int: # noqa ANN002 + if node in ECHO_NODE_TO_VIZ_NODE: + return SIZE_BY_NODE_TYPE[ECHO_NODE_TO_VIZ_NODE[node]] + return DEFAULT_NODE_SIZE + + +def echo_node_text(node: str, *args) -> str: # noqa ANN002 + return node + + +def echo_edge_text(edge: tuple, *args) -> str: # noqa ANN002 + return str(edge) + + +class PlotlyGraph: + @staticmethod + def plot( + graph: nx.graph.Graph, + positions: dict, + node_text: Callable[[Any, Any], str] | None = lambda *args: DEFAULT_NODE_NAME, + node_size: Callable[[Any, Any], int] = lambda *args: DEFAULT_NODE_SIZE, + node_color: Callable[[Any, Any], str] = lambda *args: DEFAULT_NODE_COLOR, + node_hover_text: Callable[[Any, Any], str] | None = lambda *args: DEFAULT_NODE_HOVER_TEXT, + edge_text: Callable[[Any], str] | None = lambda *args: DEFAULT_EDGE_TEXT, + edge_hover_text: Callable[[Any], str] | None = lambda *args: DEFAULT_EDGE_HOVER_TEXT, + title: str | None = None, + add_legend: bool = False, + show_node_names: bool = False, + color_connection_point_type: bool = False, + template: str = "plotly_white", + ) -> go.Figure: + + # Build edge trace + edge_x = [] + edge_y = [] + mid_edge_x = [] + mid_edge_y = [] + _edge_hover_text = [] + _edge_text = [] + for edge in graph.edges(): + x0, y0 = positions[edge[0]] + x1, y1 = positions[edge[1]] + edge_x.extend([x0, x1, None]) + edge_y.extend([y0, y1, None]) + mid_edge_x.append((x0 + x1) / 2) + mid_edge_y.append((y0 + y1) / 2) + if edge_hover_text: + _edge_hover_text.append(edge_hover_text(edge)) + if edge_text: + _edge_text.append(edge_text(edge)) + + edge_trace = go.Scatter( + x=edge_x, + y=edge_y, + line=dict(width=0.5, color="#888"), + hoverinfo="none", + mode="lines", + showlegend=False, + ) + + mid_edge_trace = go.Scatter( + x=mid_edge_x, + y=mid_edge_y, + mode="markers+text" if _edge_text else "markers", + hoverinfo="text" if _edge_hover_text else "none", + hovertext=_edge_hover_text, + text=_edge_text, + showlegend=False, + # marker=dict(size=DEFAULT_NODE_SIZE, opacity=0), + marker=dict(size=DEFAULT_NODE_SIZE, color=DEFAULT_NODE_COLOR), + ) + + # Build node trace + node_x = [] + node_y = [] + _node_text = [] + node_sizes = [] + node_colors = [] + _node_hover_text = [] + for node, val in graph.nodes.items(): + x, y = positions[node] + node_x.append(x) + node_y.append(y) + if node_text: + _node_text.append(node_text(node, val)) + if node_hover_text: + _node_hover_text.append(node_hover_text(node, val)) + node_colors.append(node_color(node, val)) + node_sizes.append(node_size(node, val)) + node_trace = go.Scatter( + x=node_x, + y=node_y, + mode="markers+text" if _node_text else "markers", + hoverinfo="text" if _node_hover_text else "none", + hovertext=_node_hover_text, + text=_node_text, + showlegend=False, + marker=dict(reversescale=True, color=node_colors, size=node_sizes, line_width=2), + ) + + trace_data = [edge_trace, mid_edge_trace, node_trace] + fig = go.Figure( + data=trace_data, + layout=go.Layout( + title=graph.name, + titlefont_size=28, + showlegend=True, + hovermode="closest", + margin=dict(b=20, l=5, r=5, t=40), + xaxis=dict(showgrid=False, zeroline=False, showticklabels=False), + yaxis=dict(showgrid=False, zeroline=False, showticklabels=False), + template=template, + ), + ) + if title: + fig.update_layout(title=dict(text=title, font=dict(size=20))) + return fig diff --git a/tests/test_objective_functions.py b/tests/test_objective_functions.py index 7f85c0e0..b933ba7c 100644 --- a/tests/test_objective_functions.py +++ b/tests/test_objective_functions.py @@ -1,3 +1,5 @@ +import pytest + from echo.configuration import Units from echo.models.agnostic import FlexPort, TellegenNode from echo.models.base import Node, OptimisationGraph @@ -106,4 +108,4 @@ def test_objectives_sum_correctly(): dt = optimise_results.get_single_objective_total_value(demand_tariff) total = optimise_results.get_total_objective_value() - assert tp + it + et + dt == total + assert tp + it + et + dt == pytest.approx(total, rel=1e-12) diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py new file mode 100644 index 00000000..356f7217 --- /dev/null +++ b/tests/test_visualisation.py @@ -0,0 +1,16 @@ +import networkx as nx +import plotly.graph_objects as go + +from echo.visualization import PlotlyGraph + + +def test_plotlygraph_plot(): + # Arrange + graph = nx.balanced_tree(3, 3) + positions = nx.spring_layout(graph) + + # Act + figure = PlotlyGraph.plot(graph=graph, positions=positions) + + # Assert + assert isinstance(figure, go.Figure)