Skip to content

Commit

Permalink
Update electronic structure results plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
edan-bainglass committed Nov 14, 2024
1 parent 84cff4f commit a65beed
Show file tree
Hide file tree
Showing 13 changed files with 121 additions and 96 deletions.
5 changes: 4 additions & 1 deletion src/aiidalab_qe/app/result/structure/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,5 +4,8 @@
class StructureResultsModel(ResultsModel):
identifier = "structure"

include = True
_this_process_label = "PwRelaxWorkChain"

@property
def include(self):
return "relax" in self.properties
4 changes: 3 additions & 1 deletion src/aiidalab_qe/app/result/summary/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,9 @@
class WorkChainSummaryModel(ResultsModel):
identifier = "summary"

include = True
@property
def include(self):
return True

def generate_report_html(self):
"""Read from the bulider parameters and generate a html for reporting
Expand Down
9 changes: 1 addition & 8 deletions src/aiidalab_qe/app/result/viewer/viewer.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def __init__(self, node: orm.Node, model: WorkChainViewerModel, **kwargs):

self._fetch_plugin_results()

# HACK
# HACK should be called from result step - fix!
self.render()

def render(self):
Expand Down Expand Up @@ -114,9 +114,6 @@ def _add_structure_panel(self):

def _fetch_plugin_results(self):
entries = get_entry_items("aiidalab_qe.properties", "result")
needs_electronic_structure = all(
identifier in self._model.properties for identifier in ("bands", "pdos")
)
for identifier, entry in entries.items():
for key in ("panel", "model"):
if key not in entry:
Expand All @@ -126,10 +123,6 @@ def _fetch_plugin_results(self):
panel = entry["panel"]
model = entry["model"]()
self._model.add_model(identifier, model)
if identifier == "electronic_structure" and needs_electronic_structure:
model.include = True
else:
model.include = identifier in self._model.properties
self.results[identifier] = panel(
identifier=identifier,
model=model,
Expand Down
2 changes: 0 additions & 2 deletions src/aiidalab_qe/common/bands_pdos/bandpdoswidget.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,15 +108,13 @@ def render(self):
description="Apply selection",
icon="pencil",
button_style="primary",
disabled=False,
)
self.update_plot_button.on_click(self._update_plot)

self.download_button = ipw.Button(
description="Download Data",
icon="download",
button_style="primary",
disabled=False,
layout=ipw.Layout(visibility="hidden"),
)
self.download_button.on_click(self._model.download_data)
Expand Down
37 changes: 5 additions & 32 deletions src/aiidalab_qe/common/panel.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,6 @@ class ResultsModel(Model, HasProcess):
this_process_uuid = tl.Unicode(allow_none=True)
process_state_notification = tl.Unicode("")

include = False
_this_process_label = ""

CSS_MAP = {
Expand All @@ -201,6 +200,10 @@ class ResultsModel(Model, HasProcess):
"created": "info",
}

@property
def include(self):
return self.identifier in self.properties

@property
def has_results(self):
return self.identifier in self.outputs
Expand All @@ -217,34 +220,6 @@ def update_process_state_notification(self):
</div>
"""

@tl.observe("process_uuid")
def _on_process_uuid_change(self, _):
super()._on_process_uuid_change(_)
if self.identifier != "summary":
self._set_this_process_uuid()

def _set_this_process_uuid(self):
if not self.process_uuid:
return
try:
self.this_process_uuid = (
orm.QueryBuilder()
.append(
orm.WorkChainNode,
filters={"uuid": self.process_node.uuid},
tag="root_process",
)
.append(
orm.WorkChainNode,
filters={"attributes.process_label": self._this_process_label},
project="uuid",
with_incoming="root_process",
)
.first(flat=True)
)
except Exception:
self.this_process_uuid = None

def _fetch_this_process_node(self):
return (
orm.QueryBuilder()
Expand Down Expand Up @@ -282,7 +257,7 @@ class ResultsPanel(Panel, t.Generic[RM]):
def __init__(self, model: RM, **kwargs):
from aiidalab_qe.common.widgets import LoadingWidget

self.loading_message = LoadingWidget(f"Loading {self.identifier} results")
self.loading_message = LoadingWidget(f"Loading {self.title.lower()} results")

super().__init__(
children=[self.loading_message],
Expand Down Expand Up @@ -338,8 +313,6 @@ def render(self):

self.rendered = True

self._model.update_process_state_notification()

def _on_load_results_click(self, _):
self.children = [self.loading_message]
self._update_view()
Expand Down
18 changes: 5 additions & 13 deletions src/aiidalab_qe/plugins/bands/result/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,25 +10,17 @@ def get_bands_node(self):
# Check for 'bands' or 'bands_projwfc' attributes within 'bands' output
if self._has_bands:
return self.outputs.bands.bands
elif self._has_bands_projwfc:
elif self._has_band_projections:
return self.outputs.bands.bands_projwfc
elif self._has_bands_output:
elif self.has_results:
# If neither 'bands' nor 'bands_projwfc' exist, use 'bands_output' itself
# This is the case for compatibility with older versions of the plugin
return self.outputs.bands

@property
def _has_bands_output(self):
return hasattr(self.outputs, "bands")

@property
def _has_bands(self):
if not self._has_bands_output:
return False
return hasattr(self.outputs.bands, "bands")
return self.has_results and hasattr(self.outputs.bands, "bands")

@property
def _has_bands_projwfc(self):
if not self._has_bands_output:
return False
return hasattr(self.outputs.bands, "bands_projwfc")
def _has_band_projections(self):
return self.has_results and hasattr(self.outputs.bands, "bands_projwfc")
6 changes: 3 additions & 3 deletions src/aiidalab_qe/plugins/electronic_structure/__init__.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
from .result import ElectronicStructureResult, ElectronicStructureResultModel
from .result import ElectronicStructureResults, ElectronicStructureResultsModel

electronic_structure = {
"result": {
"panel": ElectronicStructureResult,
"model": ElectronicStructureResultModel,
"panel": ElectronicStructureResults,
"model": ElectronicStructureResultsModel,
},
}
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from .model import ElectronicStructureResultModel
from .result import ElectronicStructureResult
from .model import ElectronicStructureResultsModel
from .result import ElectronicStructureResults

__all__ = [
"ElectronicStructureResultModel",
"ElectronicStructureResult",
"ElectronicStructureResultsModel",
"ElectronicStructureResults",
]
84 changes: 75 additions & 9 deletions src/aiidalab_qe/plugins/electronic_structure/result/model.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
from aiida import orm
from aiidalab_qe.common.panel import ResultsModel


class ElectronicStructureResultModel(ResultsModel):
# TODO if combined, this model should extend `HasModels`, and effectively
# TODO reduce to a container of Bands and PDOS, similar to its results panel
class ElectronicStructureResultsModel(ResultsModel):
identifier = "electronic_structure"

def get_pdos_node(self):
try:
return self.outputs.pdos
Expand All @@ -12,25 +17,86 @@ def get_bands_node(self):
# Check for 'bands' or 'bands_projwfc' attributes within 'bands' output
if self._has_bands:
return self.outputs.bands.bands
elif self._has_bands_projwfc:
elif self._has_band_projections:
return self.outputs.bands.bands_projwfc
elif self._has_bands_output:
# If neither 'bands' nor 'bands_projwfc' exist, use 'bands_output' itself
# This is the case for compatibility with older versions of the plugin
return self.outputs.bands

@property
def include(self):
# TODO `and` should be `or` if combined
return all(identifier in self.properties for identifier in ("bands", "pdos"))

@property
def has_results(self):
# TODO first `and` should be `or` if combined
return self._has_bands and (self._has_dos and self._has_dos_projections)

@property
def process_states(self):
nodes = self._fetch_electronic_structure_process_nodes()
return [
node.process_state.value if node and node.process_state else "queued"
for node in nodes
]

def update_process_state_notification(self):
processes = ("Bands", "PDOS")
states = self.process_states
self.process_state_notification = "\n".join(
f"""
<div class="alert alert-{self.CSS_MAP.get(state, "info")}">
<b>{process} status:</b> {state.upper()}
</div>
"""
for process, state in zip(processes, states)
)

@property
def _has_bands_output(self):
return hasattr(self.outputs, "bands")

@property
def _has_bands(self):
if not self._has_bands_output:
return False
return hasattr(self.outputs.bands, "bands")
return self._has_bands_output and hasattr(self.outputs.bands, "bands")

@property
def _has_bands_projwfc(self):
if not self._has_bands_output:
return False
return hasattr(self.outputs.bands, "bands_projwfc")
def _has_band_projections(self):
return self._has_bands_output and hasattr(self.outputs.bands, "bands_projwfc")

@property
def _has_pdos_output(self):
return hasattr(self.outputs, "pdos")

@property
def _has_dos(self):
return self._has_pdos_output and hasattr(self.outputs.pdos, "dos")

@property
def _has_dos_projections(self):
return self._has_pdos_output and hasattr(self.outputs.pdos, "projwfc")

def _fetch_electronic_structure_process_nodes(self):
return (
orm.QueryBuilder()
.append(
orm.WorkChainNode,
filters={"uuid": self.process_node.uuid},
tag="root_process",
)
.append(
orm.WorkChainNode,
filters={
"attributes.process_label": {
"in": [
"BandsWorkChain",
"PdosWorkChain",
]
}
},
with_incoming="root_process",
)
.all(flat=True)
)
15 changes: 4 additions & 11 deletions src/aiidalab_qe/plugins/electronic_structure/result/result.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,26 +5,19 @@
from aiidalab_qe.common.bands_pdos import BandPdosWidget, BandsPdosModel
from aiidalab_qe.common.panel import ResultsPanel

from .model import ElectronicStructureResultModel
from .model import ElectronicStructureResultsModel


class ElectronicStructureResult(ResultsPanel[ElectronicStructureResultModel]):
class ElectronicStructureResults(ResultsPanel[ElectronicStructureResultsModel]):
title = "Electronic Structure"
identifier = "electronic_structure"
workchain_labels = ["bands", "pdos"]

def render(self):
if self.rendered:
return

pdos_node = self._model.get_pdos_node()
def _update_view(self):
bands_node = self._model.get_bands_node()

pdos_node = self._model.get_pdos_node()
model = BandsPdosModel()
widget = BandPdosWidget(model=model, bands=bands_node, pdos=pdos_node)
widget.layout = ipw.Layout(width="1000px")
widget.render()

self.children = [widget]

self.rendered = True
12 changes: 12 additions & 0 deletions src/aiidalab_qe/plugins/pdos/result/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,20 @@


class PdosResultsModel(ResultsModel):
identifier = "pdos"

_this_process_label = "PdosWorkChain"

def get_pdos_node(self):
try:
return self.outputs.pdos
except AttributeError:
return None

@property
def _has_dos(self):
return self.has_results and hasattr(self.outputs.pdos, "dos")

@property
def _has_dos_projections(self):
return self.has_results and hasattr(self.outputs.pdos, "projwfc")
9 changes: 1 addition & 8 deletions src/aiidalab_qe/plugins/pdos/result/result.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,17 +13,10 @@ class PdosResults(ResultsPanel[PdosResultsModel]):
identifier = "pdos"
workchain_labels = ["pdos"]

def render(self):
if self.rendered:
return

def _update_view(self):
pdos_node = self._model.get_pdos_node()

model = BandsPdosModel()
widget = BandPdosWidget(model=model, pdos=pdos_node)
widget.layout = ipw.Layout(width="1000px")
widget.render()

self.children = [widget]

self.rendered = True
8 changes: 4 additions & 4 deletions tests/test_plugins_electronic_structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,15 @@ def test_electronic_structure(generate_qeapp_workchain):

from aiidalab_qe.common.bands_pdos import BandPdosWidget
from aiidalab_qe.plugins.electronic_structure.result import (
ElectronicStructureResult,
ElectronicStructureResultModel,
ElectronicStructureResults,
ElectronicStructureResultsModel,
)

workchain = generate_qeapp_workchain()
# generate structure for scf calculation
model = ElectronicStructureResultModel()
model = ElectronicStructureResultsModel()
model.process_node = workchain.node
result = ElectronicStructureResult(model=model)
result = ElectronicStructureResults(model=model)
result.render()

widget = result.children[0]
Expand Down

0 comments on commit a65beed

Please sign in to comment.