Added more testing and comments

nomad-coe · Sep 20, 2024 · 64ae0ef · 64ae0ef · github-actions · Sep 20, 2024
1 parent 83b9a41
commit 64ae0ef
Show file tree

Hide file tree

Showing 3 changed files with 349 additions and 19 deletions.
diff --git a/src/nomad_simulations/schema_packages/workflow/dft_plus_tb.py b/src/nomad_simulations/schema_packages/workflow/dft_plus_tb.py
@@ -23,18 +23,15 @@
     from nomad.datamodel.datamodel import EntryArchive
     from structlog.stdlib import BoundLogger
 
-    from nomad_simulations.schema_packages.workflow import SinglePoint
-
 from nomad.datamodel.metainfo.workflow import Link, TaskReference
 from nomad.metainfo import Quantity, Reference
 
-from nomad_simulations.schema_packages.model_method import DFT, TB, ModelMethod
-from nomad_simulations.schema_packages.workflow import (
-    BeyondDFT,
-    BeyondDFTMethod,
-)
+from nomad_simulations.schema_packages.model_method import DFT, TB
+from nomad_simulations.schema_packages.workflow import BeyondDFT, BeyondDFTMethod
 from nomad_simulations.schema_packages.workflow.base_workflows import check_n_tasks
 
+from .single_point import SinglePoint
+
 
 class DFTPlusTBMethod(BeyondDFTMethod):
     """
@@ -106,7 +103,15 @@ class DFTPlusTB(BeyondDFT):
     """
 
     @check_n_tasks(n_tasks=2)
-    def link_task_inputs_outputs(self, tasks: list[TaskReference]) -> None:
+    def link_task_inputs_outputs(
+        self, tasks: list[TaskReference], logger: 'BoundLogger'
+    ) -> None:
+        if not self.inputs or not self.outputs:
+            logger.warning(
+                'The `DFTPlusTB` workflow needs to have `inputs` and `outputs` defined in order to link with the `tasks`.'
+            )
+            return None
+
         dft_task = tasks[0]
         tb_task = tasks[1]
 
@@ -144,7 +149,7 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
                     'A `DFTPlusTB` workflow must have two `SinglePoint` tasks references.'
                 )
                 return
-            if not isinstance(task.task, 'SinglePoint'):
+            if not isinstance(task.task, SinglePoint):
                 logger.error(
                     'The referenced tasks in the `DFTPlusTB` workflow must be of type `SinglePoint`.'
                 )
@@ -158,11 +163,14 @@ def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
             tasks=self.tasks,
             tasks_names=['DFT SinglePoint Task', 'TB SinglePoint Task'],
         )
-        if method_refs is not None and len(method_refs) == 2:
-            self.method = DFTPlusTBMethod(
-                dft_method_ref=method_refs[0],
-                tb_method_ref=method_refs[1],
-            )
+        if method_refs is not None:
+            method_workflow = DFTPlusTBMethod()
+            for method in method_refs:
+                if isinstance(method, DFT):
+                    method_workflow.dft_method_ref = method
+                elif isinstance(method, TB):
+                    method_workflow.tb_method_ref = method
+            self.method = method_workflow
 
         # Resolve `tasks[*].inputs` and `tasks[*].outputs`
-        self.link_task_inputs_outputs(tasks=self.tasks)
+        self.link_task_inputs_outputs(tasks=self.tasks, logger=logger)
diff --git a/tests/workflow/test_base_workflows.py b/tests/workflow/test_base_workflows.py
@@ -84,12 +84,16 @@ def test_resolve_all_outputs(self, tasks: list[Task], result: list[Outputs]):
         [
             # no task
             (None, None),
+            # only one task
             ([TaskReference()], []),
+            # two empty tasks
             ([TaskReference(), TaskReference()], []),
+            # two tasks with only empty task
             (
                 [TaskReference(task=SinglePoint()), TaskReference(task=SinglePoint())],
                 [],
             ),
+            # two tasks with task with one input ModelSystem each
             (
                 [
                     TaskReference(
@@ -101,6 +105,7 @@ def test_resolve_all_outputs(self, tasks: list[Task], result: list[Outputs]):
                 ],
                 [],
             ),
+            # two tasks with task with one input ModelSystem each and only DFT input
             (
                 [
                     TaskReference(
@@ -121,6 +126,7 @@ def test_resolve_all_outputs(self, tasks: list[Task], result: list[Outputs]):
                 ],
                 [DFT],
             ),
+            # two tasks with task with inputs for ModelSystem and DFT and TB
             (
                 [
                     TaskReference(
File	Stmts	Miss	Cover	Missing
src/nomad_simulations
__init__.py	4	2	50%	3–4
_version.py	11	2	82%	5–6
src/nomad_simulations/schema_packages
__init__.py	15	2	87%	57–59
atoms_state.py	190	21	89%	31–33, 219–222, 246, 301–302, 370–371, 373, 555, 567–568, 629–633, 648–652, 659
basis_set.py	240	28	88%	8–9, 122–133, 172–185, 208, 391–395, 417–418, 462–465, 584, 615, 617
general.py	89	8	91%	22–25, 139, 203, 313–314, 324
model_method.py	265	75	72%	28–30, 189–192, 195–202, 315, 336–357, 373–399, 402–419, 773, 784, 826–833, 871, 890, 970, 1027, 1102, 1216
model_system.py	300	23	92%	43–45, 382, 577–580, 627–634, 808–809, 1030–1034, 1040–1041, 1049–1050, 1055, 1078
numerical_settings.py	259	61	76%	30–32, 235, 237–238, 241–244, 248–249, 256–259, 268–271, 275–278, 280–283, 288–291, 297–300, 487–514, 589, 624–627, 651, 654, 699, 701–704, 708, 712, 759, 763–784, 839–840, 907
outputs.py	120	10	92%	27–28, 270–273, 313–316, 341, 343, 380, 399
physical_property.py	102	7	93%	38–40, 220, 349–351
variables.py	86	12	86%	26–28, 116, 139, 163, 185, 207, 229, 251, 274, 294
src/nomad_simulations/schema_packages/properties
band_gap.py	51	5	90%	26–28, 153–154
band_structure.py	123	25	80%	27–29, 250–283, 296, 303, 339–340, 343, 390–391, 396
energies.py	42	9	79%	25–27, 54, 75, 100, 121, 137, 152
fermi_surface.py	17	4	76%	25–27, 58
forces.py	22	6	73%	26–28, 55, 75, 98
greens_function.py	99	13	87%	25–27, 228–229, 232, 253–254, 257, 278–279, 282, 418
hopping_matrix.py	29	5	83%	25–27, 76, 112
permittivity.py	48	8	83%	25–27, 115–123
spectral_profile.py	260	128	51%	27–29, 75–78, 113–116, 217–318, 374–386, 411–414, 434, 439–442, 484–520, 544, 591–594, 610–611, 616–622
thermodynamics.py	75	27	64%	25–27, 53, 74, 90, 99, 108, 119, 128, 155, 165, 175, 190–192, 195, 211, 231–233, 236, 252, 272–274, 277
src/nomad_simulations/schema_packages/utils
utils.py	84	18	79%	27–35, 89–98, 107–108, 113, 116
src/nomad_simulations/schema_packages/workflow
base_workflows.py	53	2	96%	24–25
dft_plus_tb.py	48	3	94%	23–24, 120
single_point.py	24	2	92%	26–27
TOTAL	2670	506	81%