Adding then supercell estimator in the Settings Panel.

src/aiidalab_qe_vibroscopy/app/settings.py

@@ -15,6 +15,60 @@
 from aiidalab_qe.common.panel import Panel
 
 
+import sys
+import os
+
+from aiida.plugins import DataFactory
+
+HubbardStructureData = DataFactory("quantumespresso.hubbard_structure")
+
+# spinner for waiting time (supercell estimations)
+spinner_html = """
+<style>
+@keyframes spin {
+  0% { transform: rotate(0deg); }
+  100% { transform: rotate(360deg); }
+}
+
+.spinner {
+  display: inline-block;
+  width: 15px;
+  height: 15px;
+}
+
+.spinner div {
+  width: 100%;
+  height: 100%;
+  border: 4px solid #f3f3f3;
+  border-top: 4px solid #3498db;
+  border-radius: 50%;
+  animation: spin 1s linear infinite;
+}
+</style>
+<div class="spinner">
+  <div></div>
+</div>
+"""
+
+
+def disable_print(func):
+    def wrapper(*args, **kwargs):
+        # Save the current standard output
+        original_stdout = sys.stdout
+        # Redirect standard output to os.devnull
+        sys.stdout = open(os.devnull, "w")
+        try:
+            # Call the function
+            result = func(*args, **kwargs)
+        finally:
+            # Restore the original standard output
+            sys.stdout.close()
+            sys.stdout = original_stdout
+        return result
+
+    return wrapper
+
+
 class Setting(Panel):
     title = "Vibrational Settings"
 
@@ -103,6 +157,7 @@ def change_supercell(_=None):
             )
         for elem in [self._sc_x, self._sc_y, self._sc_z]:
             elem.observe(change_supercell, names="value")
+            elem.observe(self._estimate_supercells, names="value")
 
         self.supercell_selector = ipw.HBox(
             children=[
@@ -124,19 +179,41 @@ def change_supercell(_=None):
         self.supercell_hint_button = ipw.Button(
             description="Size hint",
             disabled=False,
-            width="500px",
+            width="100px",
             button_style="info",
         )
         # supercell hint (15A lattice params)
         self.supercell_hint_button.on_click(self._suggest_supercell)
 
+        # reset supercell
+        self.supercell_reset_button = ipw.Button(
+            description="Reset hint",
+            disabled=False,
+            layout=ipw.Layout(width="100px"),
+            button_style="warning",
+        )
+        # supercell reset reaction
+        self.supercell_reset_button.on_click(self._reset_supercell)
+
+        # Estimate the number of supercells for frozen phonons.
+        self.supercell_number_estimator = ipw.HTML(
+            description="Number of supercells to be computed:",
+            value="0",
+            style={"description_width": "initial"},
+        )
+
         ## end supercell hint.
 
         self.supercell_widget = ipw.VBox(
             [
                 self.hint_button_help,
                 ipw.HBox(
-                    [self.supercell_selector, self.supercell_hint_button],
+                    [
+                        self.supercell_selector,
+                        self.supercell_hint_button,
+                        self.supercell_reset_button,
+                        self.supercell_number_estimator,
+                    ],
                 ),
             ]
         )
@@ -145,10 +222,26 @@ def change_supercell(_=None):
 
         self.symmetry_symprec = ipw.FloatText(
             value=1e-5,
+            max=1,
+            min=1e-7,  # Ensure the value is always positive
+            step=1e-4,  # Step value of 1e-4
             description="Symmetry tolerance (symprec):",
             style={"description_width": "initial"},
             layout={"width": "300px"},
         )
+        self.symmetry_symprec.observe(self._estimate_supercells, "value")
+
+        # reset supercell
+        self.symmetry_symprec_reset_button = ipw.Button(
+            description="Reset symprec",
+            disabled=False,
+            layout=ipw.Layout(width="150px"),
+            button_style="warning",
+        )
+        # supercell reset reaction
+        self.symmetry_symprec_reset_button.on_click(
+            lambda _: setattr(self.symmetry_symprec, "value", 1e-5)
+        )
 
         self.children = [
             ipw.VBox(
@@ -174,11 +267,22 @@ def change_supercell(_=None):
                 ],
             ),
             self.supercell_widget,
-            self.symmetry_symprec,
+            ipw.HBox(
+                [
+                    self.symmetry_symprec,
+                    self.symmetry_symprec_reset_button,
+                ],
+            ),
         ]
 
         super().__init__(**kwargs)
 
+        # we define a block for the estimation of the supercell if we ask for hint,
+        # so that we call the estimator only at the end of the supercell hint generator,
+        # and now each time after the x, y, z generation (i.e., we don't lose time).
+        # see the methods below.
+        self.block = False
+
     @tl.observe("input_structure")
     def _update_input_structure(self, change):
         if self.input_structure is not None:
@@ -204,13 +308,86 @@ def _suggest_supercell(self, _=None):
             suggested_3D = 15 // np.array(s.cell.cellpar()[:3]) + 1
 
             # if disabled, it means that it is a non-periodic direction.
+            # here we manually unobserve the `_estimate_supercells`, so it is faster
+            # and only compute when all the three directions are updated
+            self.block = True
             for direction, suggested, original in zip(
                 [self._sc_x, self._sc_y, self._sc_z], suggested_3D, s.cell.cellpar()[:3]
             ):
                 direction.value = suggested if not direction.disabled else 1
+            self.block = False
+            self._estimate_supercells()
         else:
             return
 
+    @tl.observe("input_structure")
+    @disable_print
+    def _estimate_supercells(self, _=None):
+        """_summary_
+
+        Estimate the number of supercells to be computed for frozen phonon calculation.
+        """
+        if self.block:
+            return
+
+        symprec_value = self.symmetry_symprec.value
+
+        self.symmetry_symprec.value = max(1e-5, min(symprec_value, 1))
+
+        self.supercell_number_estimator.value = spinner_html
+
+        from aiida_phonopy.data.preprocess import PreProcessData
+
+        if self.input_structure:
+            preprocess_data = PreProcessData(
+                structure=self.input_structure,
+                supercell_matrix=[
+                    [self._sc_x.value, 0, 0],
+                    [0, self._sc_y.value, 0],
+                    [0, 0, self._sc_z.value],
+                ],
+                symprec=self.symmetry_symprec.value,
+                distinguish_kinds=False,
+                is_symmetry=True,
+            )
+
+            supercells = preprocess_data.get_supercells_with_displacements()
+
+            # for now, we comment the following part, as the HubbardSD is generated in the submission step.
+            """if isinstance(self.input_structure, HubbardStructureData):
+                from aiida_vibroscopy.calculations.spectra_utils import get_supercells_for_hubbard
+                from aiida_vibroscopy.workflows.phonons.base import get_supercell_hubbard_structure
+                supercell = get_supercell_hubbard_structure(
+                    self.input_structure,
+                    self.input_structure,
+                    metadata={"store_provenance": False},
+                )
+                supercells = get_supercells_for_hubbard(
+                    preprocess_data=preprocess_data,
+                    ref_structure=supercell,
+                    metadata={"store_provenance": False},
+                )
+
+            else:
+                supercells = preprocess_data.get_supercells_with_displacements()
+            """
+            self.supercell_number_estimator.value = f"{len(supercells)}"
+
+        return
+
+    def _reset_supercell(self, _=None):
+        if self.input_structure is not None:
+            reset_supercell = []
+            self.block = True
+            for direction, periodic in zip(
+                [self._sc_x, self._sc_y, self._sc_z], self.input_structure.pbc
+            ):
+                reset_supercell.append(2 if periodic else 1)
+            (self._sc_x.value, self._sc_y.value, self._sc_z.value) = reset_supercell
+            self.block = False
+            self._estimate_supercells()
+        return
+
     def get_panel_value(self):
         """Return a dictionary with the input parameters for the plugin."""
         return {

tests/test_settings.py

@@ -75,3 +75,51 @@ def test_x_settings(generate_structure_data):
     configure_step.workchain_settings.properties["vibronic"].run.value = True
     parameters = configure_step.settings["vibronic"].get_panel_value()
     assert parameters["supercell_selector"] == [2, 1, 1]
+
+
+@pytest.mark.usefixtures("sssp")
+def test_supercell_number_estimator(generate_structure_data):
+    """Test the settings of the vibroscopy app."""
+
+    from aiidalab_qe.app.configuration import ConfigureQeAppWorkChainStep
+
+    configure_step = ConfigureQeAppWorkChainStep()
+    structure = generate_structure_data("silicon")
+    configure_step.input_structure = structure
+    configure_step.workchain_settings.properties["vibronic"].run.value = True
+    parameters = configure_step.settings["vibronic"].get_panel_value()
+    assert parameters["supercell_selector"] == [2, 2, 2]
+    assert configure_step.settings["vibronic"].supercell_number_estimator.value == "1"
+    configure_step.settings["vibronic"]._suggest_supercell()
+    parameters = configure_step.settings["vibronic"].get_panel_value()
+    assert parameters["supercell_selector"] == [4, 4, 4]
+    assert configure_step.settings["vibronic"].supercell_number_estimator.value == "1"
+    configure_step.settings["vibronic"]._sc_x.value = 3
+    configure_step.settings["vibronic"]._sc_y.value = 2
+    configure_step.settings["vibronic"]._sc_z.value = 2
+    assert configure_step.settings["vibronic"].supercell_number_estimator.value == "4"
+    configure_step.settings["vibronic"]._reset_supercell()
+    configure_step.settings["vibronic"]._sc_x.value = 2
+    configure_step.settings["vibronic"]._sc_y.value = 2
+    configure_step.settings["vibronic"]._sc_z.value = 2
+    assert configure_step.settings["vibronic"].supercell_number_estimator.value == "1"
+
+
+@pytest.mark.usefixtures("sssp")
+def test_symprec(generate_structure_data):
+    """Test the settings of the vibroscopy app."""
+
+    from aiidalab_qe.app.configuration import ConfigureQeAppWorkChainStep
+
+    configure_step = ConfigureQeAppWorkChainStep()
+    structure = generate_structure_data("silicon")
+    configure_step.input_structure = structure
+    configure_step.workchain_settings.properties["vibronic"].run.value = True
+    parameters = configure_step.settings["vibronic"].get_panel_value()
+    assert parameters["symmetry_symprec"] == 1e-5
+    configure_step.settings["vibronic"].symmetry_symprec.value = 1
+    parameters = configure_step.settings["vibronic"].get_panel_value()
+    assert parameters["symmetry_symprec"] == 1
+    configure_step.settings["vibronic"].reset()
+    parameters = configure_step.settings["vibronic"].get_panel_value()
+    assert parameters["symmetry_symprec"] == 1e-5