From 8da700446565a59affceb0dd9a0994ef118dd08a Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 May 2023 13:27:50 +0100
Subject: [PATCH] Refactor: simplify geo_opt builder param to boolean (#182)

* Refactor: simplify geo_opt builder param to boolean
---
 aiidalab_ispg/app/atmospec_steps.py | 25 ++++++++++---------------
 aiidalab_ispg/app/input_widgets.py  | 10 +++++-----
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/aiidalab_ispg/app/atmospec_steps.py b/aiidalab_ispg/app/atmospec_steps.py
index 3ad7281..b3de462 100644
--- a/aiidalab_ispg/app/atmospec_steps.py
+++ b/aiidalab_ispg/app/atmospec_steps.py
@@ -34,7 +34,7 @@
 
 @dataclass(frozen=True)
 class AtmospecParameters(OptimizationParameters):
-    geo_opt_type: str
+    optimize: bool
     excited_method: ExcitedStateMethod
     nstates: int
     es_basis: str
@@ -50,7 +50,7 @@ class AtmospecParameters(OptimizationParameters):
     method="wB97X-D4",
     basis="def2-SVP",
     solvent="None",
-    geo_opt_type="OPT",
+    optimize=True,
     excited_method=ExcitedStateMethod.TDA,
     nstates=3,
     es_basis="def2-SVP",
@@ -70,7 +70,7 @@ def __init__(self, **kwargs):
         self.molecule_settings.multiplicity.disabled = True
 
         self.geometry_settings = MolecularGeometrySettings()
-        self.geometry_settings.geo_opt_type.observe(self._observe_geo_opt_type, "value")
+        self.geometry_settings.optimize.observe(self._observe_optimize, "value")
 
         self.ground_state_settings = GroundStateSettings()
         self.ground_state_settings.method.observe(self._observe_gs_method, "value")
@@ -126,12 +126,12 @@ def _validate_input_parameters(self) -> bool:
             return False
         return True
 
-    def _observe_geo_opt_type(self, change):
+    def _observe_optimize(self, change):
         # If we don't optimize the molecule, we cannot do Wigner sampling
-        if change["new"] == "OPT":
-            self.wigner_settings.disabled = False
-        else:
+        if not change["new"]:
             self.wigner_settings.disabled = True
+        else:
+            self.wigner_settings.disabled = False
 
     def _observe_gs_sync(self, change):
         if change["new"]:
@@ -162,7 +162,7 @@ def _update_ui_from_parameters(self, parameters: AtmospecParameters) -> None:
         This function is called when we load an already finished workflow,
         and we want the input widgets to be updated accordingly
         """
-        self.geometry_settings.geo_opt_type.value = parameters.geo_opt_type
+        self.geometry_settings.optimize.value = parameters.optimize
         self.molecule_settings.charge.value = parameters.charge
         self.molecule_settings.multiplicity.value = parameters.multiplicity
         self.molecule_settings.solvent.value = parameters.solvent
@@ -187,7 +187,7 @@ def _update_ui_from_parameters(self, parameters: AtmospecParameters) -> None:
     def _get_parameters_from_ui(self) -> AtmospecParameters:
         """Prepare builder parameters from the UI input widgets"""
         return AtmospecParameters(
-            geo_opt_type=self.geometry_settings.geo_opt_type.value,
+            optimize=self.geometry_settings.optimize.value,
             charge=self.molecule_settings.charge.value,
             multiplicity=self.molecule_settings.multiplicity.value,
             solvent=self.molecule_settings.solvent.value,
@@ -340,6 +340,7 @@ def submit(self, _=None):
                 f"Excited method {bp.excited_method} not implemented"
             )
 
+        builder.optimize = bp.optimize
         builder.opt.orca.parameters = gs_opt_parameters
         builder.exc.orca.parameters = es_parameters
 
@@ -373,12 +374,6 @@ def submit(self, _=None):
         builder.exc.orca.metadata.description = "ORCA TDDFT calculation"
         builder.opt.orca.metadata.description = "ORCA geometry optimization"
 
-        # TODO: Change bp.geo_opt_type to  bp.optimize and use directly
-        if bp.geo_opt_type == "NONE":
-            builder.optimize = False
-        else:
-            builder.optimize = True
-
         # Wigner will be sampled only when optimize == True
         builder.nwigner = bp.nwigner
         builder.wigner_low_freq_thr = bp.wigner_low_freq_thr
diff --git a/aiidalab_ispg/app/input_widgets.py b/aiidalab_ispg/app/input_widgets.py
index 9030add..a591c5a 100644
--- a/aiidalab_ispg/app/input_widgets.py
+++ b/aiidalab_ispg/app/input_widgets.py
@@ -55,14 +55,14 @@ class MolecularGeometrySettings(ipw.VBox):
 
     def __init__(self, **kwargs):
         # Whether to optimize the molecule or not.
-        self.geo_opt_type = ipw.ToggleButtons(
+        self.optimize = ipw.ToggleButtons(
             options=[
-                ("Geometry as is", "NONE"),
-                ("Optimize geometry", "OPT"),
+                ("Geometry as is", False),
+                ("Optimize geometry", True),
             ],
-            value="OPT",
+            value=True,
         )
-        super().__init__(children=[self.title, self.geo_opt_type])
+        super().__init__(children=[self.title, self.optimize])
 
 
 class MoleculeSettings(ipw.VBox):