Neb workflow (#22)

Author: BTSpeake

src/aiidalab_chemshell/common/chemshell.py

@@ -30,7 +30,7 @@ class WorkflowOptions(Enum):
     GEOMETRY = 0
     SINGLE_POINT = auto()
     ATOMIC_ENERGIES = auto()
-    # NEB = auto()
+    NEB = auto()
 
     @property
     def label(self) -> str:
@@ -42,8 +42,8 @@ def label(self) -> str:
                 return "Single Point Energy"
             case WorkflowOptions.ATOMIC_ENERGIES:
                 return "Isolated Atomic Energies"
-            # case WorkflowOptions.NEB:
-            #     return "Nudged Elastic Band"
+            case WorkflowOptions.NEB:
+                return "Nudged Elastic Band"
             case _:
                 return ""
 
@@ -57,7 +57,7 @@ def tab_label(self) -> str:
                 return "SP Energy"
             case WorkflowOptions.ATOMIC_ENERGIES:
                 return "Atomic Energies"
-            # case WorkflowOptions.NEB:
-            #     return "NEB"
+            case WorkflowOptions.NEB:
+                return "NEB"
             case _:
                 return "ChemShell"

src/aiidalab_chemshell/common/structure_uploader.py

@@ -0,0 +1,121 @@
+"""Widget for selecting an input structure from various sources."""
+
+from aiida.orm import SinglefileData, StructureData, TrajectoryData
+from aiidalab_widgets_base import SmilesWidget
+from ase import Atoms
+from ipywidgets import HTML, Tab, VBox, dlink
+from traitlets import HasTraits, Instance
+
+from aiidalab_chemshell.common.database import AiiDADatabaseWidget
+from aiidalab_chemshell.common.file_handling import FileUploadWidget
+from aiidalab_chemshell.common.structure_viewer import StructureViewWidget
+
+
+class StructureSelectionWidget(VBox, HasTraits):
+    """Widget for selecting an input structre from various sources."""
+
+    structure_data = Instance(StructureData, allow_none=True)
+    structure_file = Instance(SinglefileData, allow_none=True)
+    trajectory_data = Instance(TrajectoryData, allow_none=True)
+
+    def __init__(self, **kwargs):
+        """StructureSelectionWidget constructor."""
+        super().__init__(**kwargs)
+        # upload file
+        self.file_input_widget = VBox()
+        self.file_uploader = FileUploadWidget(description="Structure file: ")
+        self.file_input_widget.children = [
+            self.file_uploader,
+        ]
+
+        # AiiDA database
+        self.database_widget = AiiDADatabaseWidget(
+            title="AiiDA Database",
+            query=[SinglefileData, StructureData],
+        )
+
+        self.smiles_widget = SmilesWidget(title="SMILES")
+
+        self.tabs = Tab()
+        self.tabs.children = [
+            self.file_input_widget,
+            self.database_widget,
+            self.smiles_widget,
+        ]
+        for i, title in enumerate(["Upload File", "AiiDA Database", "SMILES String"]):
+            self.tabs.set_title(i, title)
+
+        self.viewer = HTML("<p>No structure found...</p>")
+
+        self.children = [self.tabs, HTML("<h2>Viewer:</h2>"), self.viewer]
+
+        self.file_uploader.observe(self._on_file_upload, "file")
+        self.database_widget.observe(self._on_database_search, "data_object")
+        self.smiles_widget.observe(self._on_smiles_generation, "structure")
+
+        dlink((self.file_uploader, "file"), (self, "structure_file"))
+
+        return
+
+    def _on_file_upload(self, change: dict) -> None:
+        """When file upload button is pressed."""
+        if change["new"] != change["old"]:
+            self.viewer = StructureViewWidget()
+            self.viewer.assign_structure_from_file(
+                self.file_uploader.file.filename,
+                self.file_uploader.file.content,
+            )
+            self._update_children()
+        return
+
+    def _on_smiles_generation(self, change: dict) -> None:
+        """When SMILES string is inputted."""
+        if change["new"] != change["old"]:
+            self._create_viewer(change["new"])
+            self.structure = StructureData(ase=change["new"])
+            if self.structure_file:
+                self.structure_file = None
+        return
+
+    def _on_database_search(self, change: dict) -> None:
+        """When data is loaded from AiiDA database."""
+        if change["new"] == change["old"]:
+            return
+        if isinstance(change["new"], SinglefileData):
+            if self.structure_data:
+                self.structure_data = None
+            self.structure_file = change["new"]
+            self._on_file_upload(change)
+        elif isinstance(change["new"], StructureData):
+            if self.structure_file:
+                self.structure_file = None
+            self.structure = change["new"]
+            self._create_viewer(change["new"]._get_object_ase())
+        else:
+            self._create_viewer(None)
+        return
+
+    def _create_viewer(self, structure: Atoms | None) -> None:
+        """Create a viewer widget with the loaded ase.Atoms structure object."""
+        if structure:
+            self.viewer = StructureViewWidget()
+            self.viewer.assign_structure_from_ase(structure)
+        else:
+            self.viewer = HTML("<p>Could not visualise structure ...</p>")
+        self._update_children()
+        return
+
+    def _update_children(self) -> None:
+        self.children = [
+            self.tabs,
+            self.viewer,
+        ]
+        return
+
+    def disable(self, val: bool = True) -> None:
+        """Disable the widget and all children."""
+        for child in self.tabs.children:
+            try:
+                child.disable(True)
+            except AttributeError:
+                pass  # TODO: SmilesWidget cannot be easily disabled at present.

src/aiidalab_chemshell/models/workflow.py

@@ -11,6 +11,7 @@
 )
 
 from aiidalab_chemshell.common.chemshell import BasisSetOptions, WorkflowOptions
+from aiidalab_chemshell.models.structure import StructureInputModel
 
 
 class ChemShellWorkflowModel(HasTraits):
@@ -33,4 +34,6 @@ class ChemShellWorkflowModel(HasTraits):
     gradients = Bool(True)
     hessian = Bool(False)
 
+    structure_2 = StructureInputModel()
+
     default_guide = ""

src/aiidalab_chemshell/process.py

@@ -103,18 +103,19 @@ def submit_process(self):
                 self._submit_optimisation_workflow()
             case WorkflowOptions.ATOMIC_ENERGIES:
                 self._submit_atomic_energies_workflow()
-            case WorkflowOptions.SINGLE_POINT:
-                self._submit_core_calcjob()
             case _:
-                print("ERROR :: Invalid Workflow Specified...")
+                self._submit_core_calcjob()
         return
 
     def _submit_core_calcjob(self) -> None:
+        # Get the ChemShell code instance
         builder = load_code(self.model.resource_model.code_label).get_builder()
+        # Configure the structure input
         if self.model.structure_model.has_file:
             builder.structure = self.model.structure_model.structure_file
         else:
             builder.structure = self.model.structure_model.structure
+        # Configure the QM theory input parameters
         builder.qm_parameters = Dict(
             {
                 "theory": self.model.workflow_model.qm_theory.name,
@@ -123,6 +124,7 @@ def _submit_core_calcjob(self) -> None:
                 "basis": self.model.workflow_model.basis_set,
             }
         )
+        # Configure MM parameters if QM/MM approach specified
         if self.model.workflow_model.use_mm:
             builder.mm_parameters = Dict(
                 {
@@ -137,14 +139,25 @@ def _submit_core_calcjob(self) -> None:
                     ),
                 }
             )
-        builder.calculation_parameters = Dict(
-            {
-                "gradients": self.model.workflow_model.gradients,
-                "hessian": self.model.workflow_model.hessian,
-            }
-        )
-        if self.model.workflow_model.vibrational_analysis:
+        # Configure additional SP based tasks
+        if self.model.workflow_model.workflow == WorkflowOptions.NEB:
+            builder.optimisation_parameters = Dict({"neb": "frozen"})
+            if self.model.workflow_model.structure_2.has_file:
+                builder.structure2 = (
+                    self.model.workflow_model.structure_2.structure_file
+                )
+            else:
+                builder.structure2 = self.model.workflow_model.structure_2.structure
+        elif self.model.workflow_model.vibrational_analysis:
             builder.optimisation_parameters = Dict({"thermal": True})
+        else:
+            builder.calculation_parameters = Dict(
+                {
+                    "gradients": self.model.workflow_model.gradients,
+                    "hessian": self.model.workflow_model.hessian,
+                }
+            )
+        # Setup metadata and resource parameters
         if self.model.resource_model.ncpus > 1:
             builder.metadata.options.withmpi = True
         else:
@@ -155,6 +168,7 @@ def _submit_core_calcjob(self) -> None:
             "num_machines": 1,
             "tot_num_mpiprocs": self.model.resource_model.ncpus,
         }
+        # Submit and apply the label/description to the CalcJob
         self.node = submit(builder)
         self.node.label = self.model.resource_model.process_label
         self.node.description = self.model.resource_model.process_description

src/aiidalab_chemshell/wizards/structure.py

@@ -1,13 +1,9 @@
 """Defines the model and view components for the structure setup stage."""
 
-import ase
 import ipywidgets as ipw
-from aiida.orm import SinglefileData, StructureData
-from aiidalab_widgets_base import SmilesWidget, WizardAppWidgetStep
+from aiidalab_widgets_base import WizardAppWidgetStep
 
-from aiidalab_chemshell.common.database import AiiDADatabaseWidget
-from aiidalab_chemshell.common.file_handling import FileUploadWidget
-from aiidalab_chemshell.common.structure_viewer import StructureViewWidget
+from aiidalab_chemshell.common.structure_uploader import StructureSelectionWidget
 from aiidalab_chemshell.models.structure import StructureInputModel
 
 
@@ -42,35 +38,15 @@ def __init__(self, model: StructureInputModel, **kwargs):
             """
         )
 
-        self.tabs = ipw.Tab()
+        self.structure_uploader = StructureSelectionWidget()
 
-        # upload file
-        self.file_input_widget = ipw.VBox()
-        self.file_uploader = FileUploadWidget(description="Structure file: ")
-        self.file_input_widget.children = [
-            self.file_uploader,
-        ]
-        ipw.dlink((self.file_uploader, "file"), (self.model, "structure_file"))
-
-        # AiiDA database
-        self.database_widget = AiiDADatabaseWidget(
-            title="AiiDA Database",
-            query=[SinglefileData, StructureData],
+        ipw.dlink(
+            (self.structure_uploader, "structure_file"),
+            (self.model, "structure_file"),
+        )
+        ipw.dlink(
+            (self.structure_uploader, "structure_data"), (self.model, "structure")
         )
-
-        self.smiles_widget = SmilesWidget(title="SMILES")
-
-        self.tabs.children = [
-            self.file_input_widget,
-            self.database_widget,
-            self.smiles_widget,
-        ]
-        for i, title in enumerate(["Upload File", "AiiDA Database", "SMILES String"]):
-            self.tabs.set_title(i, title)
-
-        self.model.observe(self._on_file_upload, "structure_file")
-        self.database_widget.observe(self._on_database_search, "data_object")
-        self.smiles_widget.observe(self._on_smiles_generation, "structure")
 
     def render(self):
         """Render the wizard's contents if not already rendered."""
@@ -95,64 +71,17 @@ def render(self):
     def _update_children(self) -> None:
         self.children = [
             self.info,
-            self.tabs,
-            ipw.HTML("<h2>Viewer:</h2>"),
-            self.viewer,
+            self.structure_uploader,
             self.submit_btn,
         ]
         return
 
-    def _on_file_upload(self, change: dict) -> None:
-        """When file upload button is pressed."""
-        if self.model.has_file:
-            self.viewer = StructureViewWidget()
-            self.viewer.assign_structure_from_file(
-                self.model.structure_file.filename, self.model.structure_file.content
-            )
-            self._update_children()
-        return
-
-    def _on_smiles_generation(self, change: dict) -> None:
-        """When SMILES string is inputted."""
-        if change["new"] != change["old"]:
-            self._create_viewer(change["new"])
-            self.model.structure = StructureData(ase=change["new"])
-        return
-
-    def _on_database_search(self, change: dict) -> None:
-        """When data is loaded from AiiDA database."""
-        if change["new"] == change["old"]:
-            return
-        if isinstance(change["new"], SinglefileData):
-            self.model.structure_file = change["new"]
-            self._on_file_upload(change)
-        elif isinstance(change["new"], StructureData):
-            self.model.structure = change["new"]
-            self._create_viewer(change["new"]._get_object_ase())
-        else:
-            self._create_viewer(None)
-        return
-
-    def _create_viewer(self, structure: ase.Atoms | None) -> None:
-        """Create a viewer widget with the loaded ase.Atoms structure object."""
-        # if structure:
-        #     # self.viewer = awb.viewers.StructureDataViewer(structure=structure)
-        #     self.viewer = WeasWidget()
-        #     self.viewer.from_ase(structure)
-        # else:
-        #     self.viewer = ipw.HTML("<p>Could not visualise structure ...</p>")
-        self.viewer = StructureViewWidget()
-        self.viewer.assign_structure_from_ase(structure)
-        self._update_children()
-        return
-
     def submit_structure(self, _):
         """Submit the structure step."""
         if self.model.has_file or self.model.has_structure:
-            self.file_uploader.disable(True)
-            self.database_widget.disable(True)
             self.submit_btn.disabled = True
             self.submit_btn.description = "Submitted"
+            self.structure_uploader.disable(True)
             self.model.submitted = True
         else:
             self.model.submitted = False