Skip to content

Commit

Permalink
Adapt module and files layout for plugin implementation
Browse files Browse the repository at this point in the history
This PR is separated and prepared for #428 to avoid mixing changes, especially for moving
modules/files without changing any implementations. See #428 (comment)
for the design of restructuring the modules.
It include following changes:

- Move strucute selction to step1 module
- Move modules/folders to configuration folder (step2)
- Move modules/folders to submission (step3)
- Move modules/folders to results (steps)
- move modules/widgets to common folder
- load registred viewer in __init__ to trigger
  • Loading branch information
unkcpz committed Aug 18, 2023
1 parent 0849b9f commit 78a0535
Show file tree
Hide file tree
Showing 25 changed files with 2,017 additions and 1,951 deletions.
14 changes: 6 additions & 8 deletions qe.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -54,14 +54,12 @@
"from jinja2 import Environment\n",
"\n",
"from aiidalab_qe.app import static\n",
"from aiidalab_qe.app.process import QeAppWorkChainSelector\n",
"from aiidalab_qe.app.steps import (\n",
" ConfigureQeAppWorkChainStep,\n",
" SubmitQeAppWorkChainStep,\n",
" ViewQeAppWorkChainStatusAndResultsStep,\n",
")\n",
"from aiidalab_qe.app.structures import Examples, StructureSelectionStep\n",
"from aiidalab_qe.app.widgets import AddingTagsEditor\n",
"from aiidalab_qe.app.common.process import QeAppWorkChainSelector\n",
"from aiidalab_qe.app.common.widgets import AddingTagsEditor\n",
"from aiidalab_qe.app.configuration import ConfigureQeAppWorkChainStep\n",
"from aiidalab_qe.app.result import ViewQeAppWorkChainStatusAndResultsStep\n",
"from aiidalab_qe.app.structure import Examples, StructureSelectionStep\n",
"from aiidalab_qe.app.submission import SubmitQeAppWorkChainStep\n",
"from aiidalab_qe.version import __version__\n",
"\n",
"OptimadeQueryWidget.title = \"OPTIMADE\" # monkeypatch\n",
Expand Down
6 changes: 3 additions & 3 deletions src/aiidalab_qe/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@

import click

from aiidalab_qe.app.setup_codes import codes_are_setup
from aiidalab_qe.app.setup_codes import install as install_qe_codes
from aiidalab_qe.app.sssp import install as setup_sssp
from aiidalab_qe.app.common.setup_codes import codes_are_setup
from aiidalab_qe.app.common.setup_codes import install as install_qe_codes
from aiidalab_qe.app.submission.sssp import install as setup_sssp


@click.group()
Expand Down
9 changes: 6 additions & 3 deletions src/aiidalab_qe/app/__init__.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
"""Package for the AiiDAlab QE app."""

from .process import WorkChainSelector
from .steps import SubmitQeAppWorkChainStep, ViewQeAppWorkChainStatusAndResultsStep
from .structures import StructureSelectionStep
from .common import WorkChainSelector
from .configuration import ConfigureQeAppWorkChainStep
from .result import ViewQeAppWorkChainStatusAndResultsStep
from .structure import StructureSelectionStep
from .submission import SubmitQeAppWorkChainStep

__all__ = [
"StructureSelectionStep",
"ConfigureQeAppWorkChainStep",
"SubmitQeAppWorkChainStep",
"ViewQeAppWorkChainStatusAndResultsStep",
"WorkChainSelector",
Expand Down
7 changes: 7 additions & 0 deletions src/aiidalab_qe/app/common/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
# trigger registration of the viewer widget:
from .node_view import CalcJobNodeViewerWidget # noqa: F401
from .process import WorkChainSelector

__all__ = [
"WorkChainSelector",
]
112 changes: 112 additions & 0 deletions src/aiidalab_qe/app/common/node_view.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
"""Results view widgets (MOVE TO OTHER MODULE!)
Authors: AiiDAlab team
"""

import ipywidgets as ipw
import nglview
import traitlets as tl
from aiida import orm
from aiidalab_widgets_base import register_viewer_widget
from ase import Atoms

from .widgets import CalcJobOutputFollower, LogOutputWidget


class MinimalStructureViewer(ipw.VBox):
structure = tl.Union([tl.Instance(Atoms), tl.Instance(orm.Node)], allow_none=True)
_displayed_structure = tl.Instance(Atoms, allow_none=True, read_only=True)

background = tl.Unicode()
supercell = tl.List(tl.Int())

def __init__(self, structure, *args, **kwargs):
self._viewer = nglview.NGLWidget()
self._viewer.camera = "orthographic"
self._viewer.stage.set_parameters(mouse_preset="pymol")
ipw.link((self, "background"), (self._viewer, "background"))

self.structure = structure

super().__init__(
children=[
self._viewer,
],
*args,
**kwargs,
)

@tl.default("background")
def _default_background(self):
return "#FFFFFF"

@tl.default("supercell")
def _default_supercell(self):
return [1, 1, 1]

@tl.validate("structure")
def _valid_structure(self, change): # pylint: disable=no-self-use
"""Update structure."""
structure = change["value"]

if structure is None:
return None # if no structure provided, the rest of the code can be skipped

if isinstance(structure, Atoms):
return structure
if isinstance(structure, orm.Node):
return structure.get_ase()
raise ValueError(
"Unsupported type {}, structure must be one of the following types: "
"ASE Atoms object, AiiDA CifData or StructureData."
)

@tl.observe("structure")
def _update_displayed_structure(self, change):
"""Update displayed_structure trait after the structure trait has been modified."""
# Remove the current structure(s) from the viewer.
if change["new"] is not None:
self.set_trait("_displayed_structure", change["new"].repeat(self.supercell))
else:
self.set_trait("_displayed_structure", None)

@tl.observe("_displayed_structure")
def _update_structure_viewer(self, change):
"""Update the view if displayed_structure trait was modified."""
with self.hold_trait_notifications():
for (
comp_id
) in self._viewer._ngl_component_ids: # pylint: disable=protected-access
self._viewer.remove_component(comp_id)
self.selection = list()
if change["new"] is not None:
self._viewer.add_component(nglview.ASEStructure(change["new"]))
self._viewer.clear()
self._viewer.stage.set_parameters(clipDist=0)
self._viewer.add_representation("unitcell", diffuse="#df0587")
self._viewer.add_representation("ball+stick", aspectRatio=3.5)


class VBoxWithCaption(ipw.VBox):
def __init__(self, caption, body, *args, **kwargs):
super().__init__(children=[ipw.HTML(caption), body], *args, **kwargs)


@register_viewer_widget("process.calculation.calcjob.CalcJobNode.")
class CalcJobNodeViewerWidget(ipw.VBox):
def __init__(self, calcjob, **kwargs):
self.calcjob = calcjob
self.output_follower = CalcJobOutputFollower()
self.log_output = LogOutputWidget()

self.output_follower.calcjob_uuid = self.calcjob.uuid
self.output_follower.observe(self._observe_output_follower_lineno, ["lineno"])

super().__init__(
[ipw.HTML(f"CalcJob: {self.calcjob}"), self.log_output], **kwargs
)

def _observe_output_follower_lineno(self, _):
with self.hold_trait_notifications():
self.log_output.filename = self.output_follower.filename
self.log_output.value = "\n".join(self.output_follower.output)
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from aiida.orm import load_code
from filelock import FileLock, Timeout

from aiidalab_qe.app.widgets import ProgressBar
from aiidalab_qe.app.common.widgets import ProgressBar

__all__ = [
"QESetupWidget",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,17 +16,13 @@
import numpy as np
import traitlets
from aiida.orm import CalcJobNode, load_node
from aiidalab_widgets_base import register_viewer_widget
from aiidalab_widgets_base.utils import (
StatusHTML,
list_to_string_range,
string_range_to_list,
)
from IPython.display import HTML, Javascript, clear_output, display

# trigger registration of the viewer widget:
from aiidalab_qe.app import node_view # noqa: F401

__all__ = [
"CalcJobOutputFollower",
"LogOutputWidget",
Expand Down Expand Up @@ -350,105 +346,6 @@ def _pull_output(self):
self._output_queue.task_done()


@register_viewer_widget("process.calculation.calcjob.CalcJobNode.")
class CalcJobNodeViewerWidget(ipw.VBox):
def __init__(self, calcjob, **kwargs):
self.calcjob = calcjob
self.output_follower = CalcJobOutputFollower()
self.log_output = LogOutputWidget()

self.output_follower.calcjob_uuid = self.calcjob.uuid
self.output_follower.observe(self._observe_output_follower_lineno, ["lineno"])

super().__init__(
[ipw.HTML(f"CalcJob: {self.calcjob}"), self.log_output], **kwargs
)

def _observe_output_follower_lineno(self, _):
with self.hold_trait_notifications():
self.log_output.filename = self.output_follower.filename
self.log_output.value = "\n".join(self.output_follower.output)


class ResourceSelectionWidget(ipw.VBox):
"""Widget for the selection of compute resources."""

title = ipw.HTML(
"""<div style="padding-top: 0px; padding-bottom: 0px">
<h4>Resources</h4>
</div>"""
)
prompt = ipw.HTML(
"""<div style="line-height:120%; padding-top:0px">
<p style="padding-bottom:10px">
Specify the resources to use for the pw.x calculation.
</p></div>"""
)

def __init__(self, **kwargs):
extra = {
"style": {"description_width": "150px"},
"layout": {"min_width": "180px"},
}
self.num_nodes = ipw.BoundedIntText(
value=1, step=1, min=1, max=1000, description="Nodes", **extra
)
self.num_cpus = ipw.BoundedIntText(
value=1, step=1, min=1, description="CPUs", **extra
)

super().__init__(
children=[
self.title,
ipw.HBox(
children=[self.prompt, self.num_nodes, self.num_cpus],
layout=ipw.Layout(justify_content="space-between"),
),
]
)

def reset(self):
self.num_nodes.value = 1
self.num_cpus.value = 1


class ParallelizationSettings(ipw.VBox):
"""Widget for setting the parallelization settings."""

title = ipw.HTML(
"""<div style="padding-top: 0px; padding-bottom: 0px">
<h4>Parallelization</h4>
</div>"""
)
prompt = ipw.HTML(
"""<div style="line-height:120%; padding-top:0px">
<p style="padding-bottom:10px">
Specify the number of k-points pools for the calculations.
</p></div>"""
)

def __init__(self, **kwargs):
extra = {
"style": {"description_width": "150px"},
"layout": {"min_width": "180px"},
}
self.npools = ipw.BoundedIntText(
value=1, step=1, min=1, max=128, description="Number of k-pools", **extra
)
super().__init__(
children=[
self.title,
ipw.HBox(
children=[self.prompt, self.npools],
layout=ipw.Layout(justify_content="space-between"),
),
]
)

def reset(self):
self.npools.value = 1


class ProgressBar(ipw.HBox):
class AnimationRate(float):
pass
Expand Down
Loading

0 comments on commit 78a0535

Please sign in to comment.