From 8427c1dd2415652b10d7308d6e6db85aa79f098a Mon Sep 17 00:00:00 2001
From: jiang-yuha0 <149360440+jiang-yuha0@users.noreply.github.com>
Date: Fri, 21 Jun 2024 09:39:50 +0200
Subject: [PATCH] Outputs band_structure from the optimal W90 calculation 
 (#46)

- Output the `band_structure` from `wannier90_optimal`
---
 src/aiida_wannier90_workflows/workflows/optimize.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/aiida_wannier90_workflows/workflows/optimize.py b/src/aiida_wannier90_workflows/workflows/optimize.py
index d6b6196..3be7fe3 100644
--- a/src/aiida_wannier90_workflows/workflows/optimize.py
+++ b/src/aiida_wannier90_workflows/workflows/optimize.py
@@ -734,9 +734,6 @@ def results(self):
                     namespace="wannier90_plot",
                 )
             )
-            if "interpolated_bands" in self.outputs["wannier90_plot"]:
-                w90_bands = self.outputs["wannier90_plot"]["interpolated_bands"]
-                self.out("band_structure", w90_bands)
 
         if "optimize_reference_bands" in self.inputs:
             if self.has_run_wannier90_optimize():
@@ -744,6 +741,11 @@ def results(self):
             else:
                 # Even if I haven't run optimization, I still output bands distance if reference bands is present
                 optimal_workchain = self.ctx.workchain_wannier90
+
+            if "interpolated_bands" in optimal_workchain.outputs:
+                # Override the `band_strucure` from W90BandsWorkChain
+                w90_optimal_bands = optimal_workchain.outputs["interpolated_bands"]
+                self.out("band_structure", w90_optimal_bands)
             bandsdist = self._get_bands_distance(optimal_workchain)
             bandsdist = orm.Float(bandsdist)
             bandsdist.store()