fram-core 0.0.0__py3-none-any.whl → 0.1.0a1__py3-none-any.whl

framcore/utils/get_regional_volumes.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import contextlib
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import numpy as np
+from numpy.typing import NDArray
+
+from framcore.attributes import FlowVolume
+from framcore.components import Component, Flow, Node
+from framcore.events import send_warning_event
+from framcore.expressions import get_unit_conversion_factor
+from framcore.expressions._utils import _load_model_and_create_model_db
+from framcore.metadata import Member
+from framcore.querydbs import QueryDB
+from framcore.timeindexes import FixedFrequencyTimeIndex, SinglePeriodTimeIndex
+from framcore.utils import FlowInfo, get_flow_infos, get_node_to_commodity, get_supported_components
+
+if TYPE_CHECKING:
+    from framcore import Model
+
+
+class RegionalVolumes:
+    """
+    Container for regional energy volumes.
+
+    Stores production, consumption, import, and export vectors for each node and category.
+    Provides methods to access these aggregated results.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the RegionalVolumes instance with empty dictionaries for production, consumption, import, and export."""
+        self._production: dict[str, dict[str, NDArray]] = dict()
+        self._consumption: dict[str, dict[str, NDArray]] = dict()
+        self._export: dict[str, dict[str, NDArray]] = dict()
+        self._import: dict[str, dict[str, NDArray]] = dict()
+
+    def get_production(self) -> dict[str, dict[str, NDArray]]:
+        """Return dict with production vector by category for each node."""
+        return self._production
+
+    def get_consumption(self) -> dict[str, dict[str, NDArray]]:
+        """Return dict with consumption vector by category for each node."""
+        return self._consumption
+
+    def get_export(self) -> dict[str, dict[str, NDArray]]:
+        """Return nested dict with export vector for each trade partner to an exporting node."""
+        return self._export
+
+    def get_import(self) -> dict[str, dict[str, NDArray]]:
+        """Return nested dict with import vector for each trade partner to an importing node."""
+        return self._import
+
+
+def _get_meta_value(key: str, v: Node | Flow, category_level: str) -> str:
+    """Get member meta value from component."""
+    meta = v.get_meta(category_level)
+    if not isinstance(meta, Member):
+        message = f"Expected for key {key} metadata of type Member, got {meta}"
+        raise ValueError(message)
+    return meta.get_value()
+
+
+def _get_vector(
+    flow: Flow,
+    is_ingoing: bool,
+    commodity: str,
+    node_to_commodity: dict[str, str],
+    db: QueryDB,
+    data_period: SinglePeriodTimeIndex,
+    scenario_period: FixedFrequencyTimeIndex,
+    unit: str,
+    is_float32: bool,
+) -> FlowVolume:
+    arrows = flow.get_arrows()
+    if len(arrows) == 1:
+        volume = flow.get_volume()
+        return volume.get_scenario_vector(
+            db=db,
+            scenario_horizon=scenario_period,
+            level_period=data_period,
+            unit=unit,
+            is_float32=is_float32,
+        )
+
+    arrows = [a for a in flow.get_arrows() if a.is_ingoing() == is_ingoing and node_to_commodity[a.get_node()] == commodity]
+    if len(arrows) != 1:
+        message = f"Expected one arrow, got {arrows}"
+        raise ValueError(message)
+    arrow = arrows[0]
+
+    arrow_volumes = flow.get_arrow_volumes()
+
+    if arrow in arrow_volumes:
+        volume = arrow_volumes[arrow]
+        return volume.get_scenario_vector(
+            db=db,
+            scenario_horizon=scenario_period,
+            level_period=data_period,
+            unit=unit,
+            is_float32=is_float32,
+        )
+
+    # we have to calculate using volume and conversion
+    volume = flow.get_volume()
+
+    main_node = flow.get_main_node()
+    main_arrows = [a for a in flow.get_arrows() if a.get_node() == main_node]
+    if len(main_arrows) != 1:
+        message = f"Expected exactly one arrow connected to main node of flow. Got {main_arrows}"
+        raise ValueError(message)
+    main_arrow = main_arrows[0]
+
+    if arrow == main_arrow:
+        return volume.get_scenario_vector(
+            db=db,
+            scenario_horizon=scenario_period,
+            level_period=data_period,
+            unit=unit,
+            is_float32=is_float32,
+        )
+
+    main_units = main_arrow.get_conversion_unit_set(db)
+    if not main_units:
+        return volume.get_scenario_vector(
+            db=db,
+            scenario_horizon=scenario_period,
+            level_period=data_period,
+            unit=unit,
+            is_float32=is_float32,
+        )
+
+    # we must convert to correct unit
+    arrow_units = arrow.get_conversion_unit_set(db)
+
+    a_main_unit = next(iter(main_units))
+    a_arrow_unit = next(iter(arrow_units))
+
+    unit_conversion_factor = get_unit_conversion_factor(
+        from_unit=f"(({a_arrow_unit}) / ({a_main_unit}))",
+        to_unit=unit,
+    )
+
+    vector = volume.get_scenario_vector(
+        db=db,
+        scenario_horizon=scenario_period,
+        level_period=data_period,
+        unit=a_main_unit,
+        is_float32=is_float32,
+    )
+
+    if arrow.has_profile():
+        conversion_vector = arrow.get_scenario_vector(
+            db=db,
+            scenario_horizon=scenario_period,
+            level_period=data_period,
+            unit=a_arrow_unit,
+            is_float32=is_float32,
+        )
+        np.multiply(vector, conversion_vector, out=vector)
+        np.multiply(vector, unit_conversion_factor, out=vector)
+        return vector
+
+    conversion_value = arrow.get_data_value(
+        db=db,
+        scenario_horizon=scenario_period,
+        level_period=data_period,
+        unit=a_arrow_unit,
+    )
+    np.multiply(vector, conversion_value * unit_conversion_factor, out=vector)
+
+    return vector
+
+
+# TODO: More options: node_category, consumption_category, production_category, with_trade_partners
+
+
+def _check_category(category, flow_id, flow_info) -> None:
+    pass
+
+
+def get_regional_volumes(
+    db: Model | QueryDB,
+    commodity: str,
+    node_category: str,
+    production_category: str,
+    consumption_category: str,
+    data_period: SinglePeriodTimeIndex,
+    scenario_period: FixedFrequencyTimeIndex,
+    unit: str,
+    is_float32: bool = True,
+) -> RegionalVolumes:
+    """Calculate aggregated production, consumption, import and export."""
+    db = _load_model_and_create_model_db(db)
+
+    if not isinstance(is_float32, bool):
+        message = f"Expected bool for is_float32, got {is_float32}"
+        raise ValueError(message)
+
+    domain_components = {k: v for k, v in db.get_data().items() if isinstance(v, Component)}
+
+    graph: dict[str, Node | Flow] = get_supported_components(
+        components=domain_components,
+        supported_types=(Node, Flow),
+        forbidden_types=tuple(),
+    )
+
+    flows: dict[str, Flow] = {k: v for k, v in graph.items() if isinstance(v, Flow)}
+    nodes: dict[str, Node] = {k: v for k, v in graph.items() if isinstance(v, Node)}
+
+    node_to_commodity = get_node_to_commodity(graph)
+
+    # only nodes of prefered commodity
+    nodes_of_commodity: dict[str, Node] = {k: v for k, v in nodes.items() if v.get_commodity() == commodity}
+
+    # Mapping of node to category of prefered node level
+    node_to_category: dict[str, str] = {k: _get_meta_value(k, v, node_category) for k, v in nodes_of_commodity.items()}
+
+    category_to_nodes: dict[str, set[str]] = defaultdict(set)
+    visited = set()
+    for node_id, category in node_to_category.items():
+        assert node_id not in visited, f"{node_id} is duplicated"
+        category_to_nodes[category].add(node_id)
+        visited.add(node_id)
+
+    direct_production: dict[str, dict[str, list[Flow]]] = dict()
+    direct_consumption: dict[str, dict[str, list[Flow]]] = dict()
+    converted_production: dict[str, dict[str, list[Flow]]] = dict()
+    converted_consumption: dict[str, dict[str, list[Flow]]] = dict()
+    import_: dict[str, dict[str, list[Flow]]] = dict()
+    export: dict[str, dict[str, list[Flow]]] = dict()
+
+    for flow_id, flow in flows.items():
+        flow_infos = get_flow_infos(flow, node_to_commodity)
+
+        prod_category = None
+        cons_category = None
+        with contextlib.suppress(Exception):
+            prod_category = _get_meta_value(flow_id, flow, production_category)
+        with contextlib.suppress(Exception):
+            cons_category = _get_meta_value(flow_id, flow, consumption_category)
+
+        for flow_info in flow_infos:
+            flow_info: FlowInfo
+            if flow_info.category == "direct_in" and flow_info.commodity_in == commodity:
+                _check_category(prod_category, flow_id, flow_info)
+                node_category = node_to_category[flow_info.node_in]
+                if node_category not in direct_production:
+                    direct_production[node_category] = defaultdict(list)
+                direct_production[node_category][prod_category].append(flow)
+
+            elif flow_info.category == "conversion" and flow_info.commodity_in == commodity:
+                _check_category(prod_category, flow_id, flow_info)
+                node_category = node_to_category[flow_info.node_in]
+                if node_category not in converted_production:
+                    converted_production[node_category] = defaultdict(list)
+                converted_production[node_category][prod_category].append(flow)
+
+            elif flow_info.category == "direct_out" and flow_info.commodity_out == commodity:
+                _check_category(cons_category, flow_id, flow_info)
+                node_category = node_to_category[flow_info.node_out]
+                if node_category not in direct_consumption:
+                    direct_consumption[node_category] = defaultdict(list)
+                direct_consumption[node_category][cons_category].append(flow)
+
+            elif flow_info.category == "conversion" and flow_info.commodity_out == commodity:
+                _check_category(cons_category, flow_id, flow_info)
+                node_category = node_to_category[flow_info.node_out]
+                if node_category not in converted_consumption:
+                    converted_consumption[node_category] = defaultdict(list)
+                converted_consumption[node_category][cons_category].append(flow)
+
+            elif flow_info.category == "transport":
+                if node_to_commodity[flow_info.node_in] != commodity:
+                    continue
+                category_in = node_to_category[flow_info.node_in]
+                category_out = node_to_category[flow_info.node_out]
+                if category_in == category_out:
+                    continue
+
+                if category_in not in import_:
+                    import_[category_in] = defaultdict(list)
+                import_[category_in][category_out].append(flow)
+
+                if category_out not in export:
+                    export[category_out] = defaultdict(list)
+                export[category_out][category_in].append(flow)
+
+    num_periods = scenario_period.get_num_periods()
+    dtype = np.float32 if is_float32 else np.float64
+
+    out = RegionalVolumes()
+
+    # direct
+    for flow_dict, out_dict, is_ingoing in [(direct_production, out.get_production(), True), (direct_consumption, out.get_consumption(), False)]:
+        for node_category, flow_categories in flow_dict.items():
+            if node_category not in out_dict:
+                out_dict[node_category] = dict()
+            for flow_category, flows in flow_categories.items():
+                x = np.zeros(num_periods, dtype=dtype)
+                for flow in set(flows):
+                    try:
+                        vector = _get_vector(
+                            flow=flow,
+                            is_ingoing=is_ingoing,
+                            commodity=commodity,
+                            node_to_commodity=node_to_commodity,
+                            db=db,
+                            scenario_period=scenario_period,
+                            data_period=data_period,
+                            unit=unit,
+                            is_float32=is_float32,
+                        )
+                        np.add(x, vector, out=x)
+                    except Exception as e:
+                        send_warning_event(flow, f"Could not get direct production or consumption for flow {flow}: {e}")
+                out_dict[node_category][flow_category] = x
+
+    # converted
+    for flow_dict, out_dict, is_ingoing in [(converted_production, out.get_production(), True), (converted_consumption, out.get_consumption(), False)]:
+        for node_category, flow_categories in flow_dict.items():
+            if node_category not in out_dict:
+                out_dict[node_category] = dict()
+            for flow_category, flows in flow_categories.items():
+                x = out_dict[node_category][flow_category] if flow_category in out_dict[node_category] else np.zeros(num_periods, dtype=dtype)
+                for flow in set(flows):
+                    try:
+                        vector = _get_vector(
+                            flow=flow,
+                            is_ingoing=is_ingoing,
+                            commodity=commodity,
+                            node_to_commodity=node_to_commodity,
+                            db=db,
+                            scenario_period=scenario_period,
+                            data_period=data_period,
+                            unit=unit,
+                            is_float32=is_float32,
+                        )
+                        np.add(x, vector, out=x)
+                    except Exception as e:
+                        send_warning_event(flow, f"Could not get indirect production or consumption for flow {flow}: {e}")
+                out_dict[node_category][flow_category] = x
+
+    # trade
+    for flow_dict, out_dict, is_ingoing in [(import_, out.get_import(), True), (export, out.get_export(), False)]:
+        for category, trade_partners in flow_dict.items():
+            out_dict[category] = dict()
+            for trade_partner, flows in trade_partners.items():
+                x = np.zeros(num_periods, dtype=dtype)
+                for flow in set(flows):
+                    try:
+                        vector = _get_vector(
+                            flow=flow,
+                            is_ingoing=is_ingoing,
+                            commodity=commodity,
+                            node_to_commodity=node_to_commodity,
+                            db=db,
+                            scenario_period=scenario_period,
+                            data_period=data_period,
+                            unit=unit,
+                            is_float32=is_float32,
+                        )
+                        np.add(x, vector, out=x)
+                    except Exception as e:
+                        send_warning_event(flow, f"Could not get trade for flow {flow}: {e}")
+                out_dict[category][trade_partner] = x
+
+    return out

framcore/utils/get_supported_components.py

@@ -0,0 +1,60 @@
+from framcore.components import Component
+
+
+def get_supported_components(
+    components: dict[str, Component],
+    supported_types: tuple[type[Component]],
+    forbidden_types: tuple[type[Component]],
+) -> dict[str, Component]:
+    """Return simplified version of components in compliance with specified component types."""
+    output: dict[str, Component] = {}
+    errors: list[str] = []
+
+    _simplify_until_supported(
+        output,
+        errors,
+        components,
+        supported_types,
+        forbidden_types,
+    )
+
+    if errors:
+        message = "\n".join(errors)
+        raise ValueError(message)
+
+    return output
+
+
+def _simplify_until_supported(
+    output: dict[str, Component],
+    errors: list[str],
+    candidates: dict[str, Component],
+    supported_types: tuple[type[Component]],
+    forbidden_types: tuple[type[Component]],
+) -> None:
+    for name, component in candidates.items():
+        if isinstance(component, forbidden_types):
+            message = f"{component.get_top_parent()} has forbidden component {component}"
+            errors.append(message)
+
+        elif isinstance(component, supported_types):
+            output[name] = component
+
+        else:
+            simpler_components = component.get_simpler_components(name)
+
+            if not simpler_components:
+                message = (
+                    f"Failed to support component. Reached bottom level component {component} with top level "
+                    f"parent {component.get_top_parent()}. No component in the hierarchy was supported."
+                )
+                errors.append(message)
+
+            else:
+                _simplify_until_supported(
+                    output,
+                    errors,
+                    simpler_components,
+                    supported_types,
+                    forbidden_types,
+                )

framcore/utils/global_energy_equivalent.py

@@ -0,0 +1,46 @@
+from framcore.components import Component, HydroModule
+from framcore.curves import Curve
+from framcore.expressions import Expr
+from framcore.metadata import LevelExprMeta
+from framcore.timevectors import ConstantTimeVector, TimeVector
+
+
+def get_hydro_downstream_energy_equivalent(
+    data: dict[str, Component | TimeVector | Curve | Expr],
+    module_name: str,
+    power_node: str | None = None,
+) -> Expr:
+    """
+    Get the sum downstream energy equivalent for a hydro module.
+
+    Either count all downstream energy equivalents, or only those that are connected to the given power_node.
+    """
+    if data[module_name].get_pump() and data[module_name].get_pump().get_from_module() == module_name:  # transport pump
+        pump_power_node = data[module_name].get_pump().get_power_node()
+        pump_to = data[module_name].get_pump().get_to_module()
+        energy_equivalent = get_hydro_downstream_energy_equivalent(data, pump_to, power_node)  # continue downstream of pump_to module
+        if power_node in (pump_power_node, None):
+            return energy_equivalent - data[module_name].get_pump().get_energy_eq().get_level()  # pumps has negative energy equivalents
+        return energy_equivalent
+
+    energy_equivalent = 0
+    if data[module_name].get_generator():  # hydro generator
+        module_power_node = data[module_name].get_generator().get_power_node()
+        if power_node in (module_power_node, None):
+            energy_equivalent += data[module_name].get_generator().get_energy_eq().get_level()
+    if data[module_name].get_release_to():  # continue from release_to module
+        release_to = data[module_name].get_release_to()
+        energy_equivalent += get_hydro_downstream_energy_equivalent(data, release_to, power_node)
+    return energy_equivalent
+
+
+def set_global_energy_equivalent(data: dict[str, Component | TimeVector | Curve | Expr], metakey_energy_eq_downstream: str) -> None:
+    """Set the downstream energy equivalent of all HydroModules. Set to 1 for other types of components?."""
+    for module_name, module in data.items():
+        if isinstance(module, HydroModule) and module.get_reservoir():
+            energy_equivalent = get_hydro_downstream_energy_equivalent(data, module_name)
+            if energy_equivalent == 0:
+                message = f"HydroModule {module_name} has no downstream energy equivalents."
+                module.send_warning_event(message)
+                energy_equivalent = ConstantTimeVector(scalar=0.0, unit="kWh/m3", is_max_level=False)
+            module.add_meta(metakey_energy_eq_downstream, LevelExprMeta(energy_equivalent))

framcore/utils/isolate_subnodes.py

@@ -0,0 +1,163 @@
+"""Demo to show how we can use the core to write some functions we need."""
+
+from collections import defaultdict
+from copy import copy
+from time import time
+
+from framcore import Model
+from framcore.components import Component, Flow, Node
+from framcore.events import send_debug_event
+from framcore.utils import get_node_to_commodity, get_supported_components, is_transport_by_commodity
+
+
+def _is_boundary_flow(flow: Flow, nodes: set[str]) -> bool:
+    arrows = flow.get_arrows()
+    x, y = tuple(arrows)  # has len 2
+    return int(x.get_node() in nodes) + int(y.get_node() in nodes) == 1
+
+
+def _is_member(node: Node, meta_key: str, members: set[str]) -> bool:
+    meta = node.get_meta(meta_key)
+    value = meta.get_value()
+    return value in members
+
+
+def isolate_subnodes(model: Model, commodity: str, meta_key: str, members: list[str]) -> None:
+    """
+    Delete nodes of commodity named using meta_key except members and boundary nodes and flows.
+
+    Boudary nodes are set exogenous and all flows pointing to them except boundary flows into or out from member nodes.
+    """
+    t = time()
+
+    data = model.get_data()
+    counts_before = model.get_content_counts()
+
+    has_not_converged = True
+    num_iterations = 0
+
+    while has_not_converged:
+        num_iterations += 1
+
+        n_data_before = len(data)
+
+        # We need copy of components to set _parent None so component becomes top_parent in upcoming code
+        components: dict[str, Component] = {k: copy(v) for k, v in data.items() if isinstance(v, Component)}
+        for c in components.values():
+            c: Component
+            c._parent = None  # noqa: SLF001
+
+        node_to_commodity = get_node_to_commodity(components)
+
+        parent_keys: dict[Component, str] = {v: k for k, v in components.items()}
+
+        graph: dict[str, Node | Flow] = get_supported_components(components, (Node, Flow), tuple())
+
+        parent_to_components = defaultdict(set)
+        for c in graph.values():
+            parent_to_components[c.get_top_parent()].add(c)
+
+        nodes: dict[str, Node] = {k: v for k, v in graph.items() if isinstance(v, Node)}
+        flows: dict[str, Flow] = {k: v for k, v in graph.items() if isinstance(v, Flow)}
+
+        commodity_nodes: dict[str, Node] = {k: v for k, v in nodes.items() if commodity == v.get_commodity()}
+        for k, v in commodity_nodes.items():
+            assert v.get_meta(meta_key), f"missing meta_key {meta_key} node_id {k}"
+
+        inside_nodes: dict[str, Node] = {k: v for k, v in commodity_nodes.items() if _is_member(v, meta_key, members)}
+
+        transports: dict[str, Flow] = {k: v for k, v in flows.items() if is_transport_by_commodity(v, node_to_commodity, commodity)}
+
+        boundary_flows: dict[str, Flow] = {k: v for k, v in transports.items() if _is_boundary_flow(v, inside_nodes.keys())}
+
+        boundary_nodes: dict[str, Node] = dict()
+        for flow_id, flow in boundary_flows.items():
+            for a in flow.get_arrows():
+                node_id = a.get_node()
+                if node_id not in inside_nodes:
+                    boundary_nodes[node_id] = nodes[node_id]
+
+        outside_nodes: dict[str, Node] = {k: v for k, v in commodity_nodes.items() if not (k in inside_nodes or k in boundary_nodes)}
+
+        deletes: set[str] = set()
+
+        deletes.update(outside_nodes.keys())
+        deletes.update(boundary_nodes.keys())  # will be kept in delete step below
+        deletes.update(boundary_flows.keys())  # will be kept in delete step below
+
+        # delete flows delivering to deleted node
+        for k, flow in flows.items():
+            for a in flow.get_arrows():
+                if a.get_node() in deletes:
+                    deletes.add(k)
+                    break  # goto next k, flow
+
+        # needed for next two steps
+        node_to_flows: dict[str, set[str]] = defaultdict(set)
+        flow_to_nodes: dict[str, set[str]] = defaultdict(set)
+        for flow_id, flow in flows.items():
+            for arrow in flow.get_arrows():
+                node_id = arrow.get_node()
+                node_to_flows[node_id].add(flow_id)
+                flow_to_nodes[flow_id].add(node_id)
+
+        # delete disconnected subgraphs
+        remaining = set(n for n in nodes if n not in commodity_nodes)
+        while remaining:
+            is_disconnected_subgraph = True
+            subgraph = set()
+            possible_members = set()
+            possible_members.add(remaining.pop())
+            while possible_members:
+                member = possible_members.pop()
+                if member in subgraph:  # avoid cycle
+                    continue
+                if member in flows:
+                    subgraph.add(member)
+                    for node in flow_to_nodes[member]:
+                        if node not in outside_nodes or node not in boundary_nodes:
+                            possible_members.add(node)
+                            if node in inside_nodes:
+                                is_disconnected_subgraph = False
+                else:
+                    subgraph.add(member)
+                    for flow in node_to_flows[member]:
+                        possible_members.add(flow)
+            if is_disconnected_subgraph:
+                deletes.update(subgraph)
+
+        for key in deletes:
+            if (key in boundary_flows) or (key in boundary_nodes):
+                continue
+
+            if key not in graph:
+                continue
+
+            parent_key = parent_keys[graph[key].get_top_parent()]
+
+            if parent_key in data:
+                del data[parent_key]
+
+        n_data_after = len(data)
+
+        if n_data_after == n_data_before:
+            has_not_converged = False
+
+    counts_after = model.get_content_counts()
+
+    added_components = counts_after["components"] - counts_before["components"]
+    if added_components.total() > 0:
+        message = f"Expected only deleted components. Got additions {added_components}"
+        raise RuntimeError(message)
+
+    deleted_components = counts_before["components"] - counts_after["components"]
+
+    for node_id in boundary_nodes:
+        if node_id in data:
+            node: Node = data[node_id]
+            node.set_exogenous()
+            if not node.get_price().has_level():
+                message = f"{node_id} set to be exogenous, but no price is available."
+                raise RuntimeError(message)
+
+    send_debug_event(isolate_subnodes, f"Used {num_iterations} iterations and {round(time() - t, 2)} seconds and deleted {deleted_components}")