Use the new included_in_stat hierarchy in the Energy Sankey card (#25306)

src/components/chart/ha-sankey-chart.ts

@@ -105,10 +105,41 @@ export class HaSankeyChart extends LitElement {
 
   private _createData = memoizeOne((data: SankeyChartData, width = 0) => {
     const filteredNodes = data.nodes.filter((n) => n.value > 0);
-    const indexes = [...new Set(filteredNodes.map((n) => n.index))];
+    const indexes = [...new Set(filteredNodes.map((n) => n.index))].sort();
     const depthMap = new Map<number, number>();
-    indexes.sort().forEach((index, i) => {
+    const sections: Node[][] = [];
+    indexes.forEach((index, i) => {
       depthMap.set(index, i);
+      const nodesWithIndex = filteredNodes.filter((n) => n.index === index);
+      if (nodesWithIndex.length > 0) {
+        sections.push(
+          sections.length > 0
+            ? nodesWithIndex.sort((a, b) => {
+                // sort by the order of their parents in the previous section with orphans at the end
+                const aParentIndex = this._findParentIndex(
+                  a.id,
+                  data.links,
+                  sections
+                );
+                const bParentIndex = this._findParentIndex(
+                  b.id,
+                  data.links,
+                  sections
+                );
+                if (aParentIndex === bParentIndex) {
+                  return 0;
+                }
+                if (aParentIndex === -1) {
+                  return 1;
+                }
+                if (bParentIndex === -1) {
+                  return -1;
+                }
+                return aParentIndex - bParentIndex;
+              })
+            : nodesWithIndex
+        );
+      }
     });
     const links = this._processLinks(filteredNodes, data.links);
     const sectionWidth = width / indexes.length;
@@ -117,7 +148,7 @@ export class HaSankeyChart extends LitElement {
     return {
       id: "sankey",
       type: "sankey",
-      nodes: filteredNodes.map((node) => ({
+      nodes: sections.flat().map((node) => ({
         id: node.id,
         value: node.value,
         itemStyle: {
@@ -227,6 +258,23 @@ export class HaSankeyChart extends LitElement {
     return links;
   }
 
+  private _findParentIndex(id: string, links: Link[], sections: Node[][]) {
+    const parent = links.find((l) => l.target === id)?.source;
+    if (!parent) {
+      return -1;
+    }
+    let offset = 0;
+    for (let i = sections.length - 1; i >= 0; i--) {
+      const section = sections[i];
+      const index = section.findIndex((n) => n.id === parent);
+      if (index !== -1) {
+        return offset + index;
+      }
+      offset += section.length;
+    }
+    return -1;
+  }
+
   static styles = css`
     :host {
       display: block;

src/panels/lovelace/cards/energy/hui-energy-sankey-card.ts

@@ -227,6 +227,7 @@ class HuiEnergySankeyCard
 
     let untrackedConsumption = homeNode.value;
     const deviceNodes: Node[] = [];
+    const parentLinks: Record<string, string> = {};
     prefs.device_consumption.forEach((device, idx) => {
       const value =
         device.stat_consumption in this._data!.stats
@@ -238,7 +239,7 @@ class HuiEnergySankeyCard
         return;
       }
       untrackedConsumption -= value;
-      deviceNodes.push({
+      const node = {
         id: device.stat_consumption,
         label:
           device.name ||
@@ -251,12 +252,24 @@ class HuiEnergySankeyCard
         tooltip: `${formatNumber(value, this.hass.locale)} kWh`,
         color: getGraphColorByIndex(idx, computedStyle),
         index: 4,
+        parent: device.included_in_stat,
+      };
+      if (node.parent) {
+        parentLinks[node.id] = node.parent;
+        links.push({
+          source: node.parent,
+          target: node.id,
         });
+      }
+      deviceNodes.push(node);
     });
+    const devicesWithoutParent = deviceNodes.filter(
+      (node) => !parentLinks[node.id]
+    );
 
     const { group_by_area, group_by_floor } = this._config;
     if (group_by_area || group_by_floor) {
-      const { areas, floors } = this._groupByFloorAndArea(deviceNodes);
+      const { areas, floors } = this._groupByFloorAndArea(devicesWithoutParent);
 
       Object.keys(floors)
         .sort(
@@ -310,7 +323,6 @@ class HuiEnergySankeyCard
 
             // Link devices to the appropriate target (area, floor, or home)
             areas[areaId].devices.forEach((device) => {
-              nodes.push(device);
               links.push({
                 source: targetNodeId,
                 target: device.id,
@@ -320,8 +332,7 @@ class HuiEnergySankeyCard
           });
         });
     } else {
-      deviceNodes.forEach((deviceNode) => {
+      devicesWithoutParent.forEach((deviceNode) => {
-        nodes.push(deviceNode);
         links.push({
           source: "home",
           target: deviceNode.id,
@@ -329,6 +340,12 @@ class HuiEnergySankeyCard
         });
       });
     }
+    const deviceSections = this._getDeviceSections(parentLinks, deviceNodes);
+    deviceSections.forEach((section, index) => {
+      section.forEach((node: Node) => {
+        nodes.push({ ...node, index: 4 + index });
+      });
+    });
 
     // untracked consumption
     if (untrackedConsumption > 0) {
@@ -340,7 +357,7 @@ class HuiEnergySankeyCard
         value: untrackedConsumption,
         tooltip: `${formatNumber(untrackedConsumption, this.hass.locale)} kWh`,
         color: computedStyle.getPropertyValue("--state-unavailable-color"),
-        index: 4,
+        index: 3 + deviceSections.length,
       });
       links.push({
         source: "home",
@@ -428,6 +445,31 @@ class HuiEnergySankeyCard
     return { areas, floors };
   }
 
+  protected _getDeviceSections(
+    parentLinks: Record<string, string>,
+    deviceNodes: Node[]
+  ): Node[][] {
+    const parentSection: Node[] = [];
+    const childSection: Node[] = [];
+    const parentIds = Object.values(parentLinks);
+    const remainingLinks: typeof parentLinks = {};
+    deviceNodes.forEach((deviceNode) => {
+      if (parentIds.includes(deviceNode.id)) {
+        parentSection.push(deviceNode);
+        remainingLinks[deviceNode.id] = parentLinks[deviceNode.id];
+      } else {
+        childSection.push(deviceNode);
+      }
+    });
+    if (parentSection.length > 0) {
+      return [
+        ...this._getDeviceSections(remainingLinks, parentSection),
+        childSection,
+      ];
+    }
+    return [deviceNodes];
+  }
+
   static styles = css`
     :host {
       display: block;