Fix types

edan-bainglass · Aug 12, 2024 · 0419ef4 · 0419ef4
1 parent a41c5fb
commit 0419ef4
Show file tree

Hide file tree

Showing 9 changed files with 109 additions and 94 deletions.
diff --git a/src/aiida_quantum_transport/calculations/base.py b/src/aiida_quantum_transport/calculations/base.py
@@ -1,20 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from aiida.engine import CalcJob
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class BaseCalculation(CalcJob):
     """docstring"""
 
     _default_parser_name = ""
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)

diff --git a/src/aiida_quantum_transport/calculations/current.py b/src/aiida_quantum_transport/calculations/current.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class CurrentCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.current"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -114,12 +110,24 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             transmission_folder_path,
         ]
 
-        hybridization_data: orm.RemoteData = (
-            self.inputs.hybridization.remote_results_folder
-        )
-        transmission_data: orm.RemoteData = (
-            self.inputs.transmission.remote_results_folder
-        )
+        hybridization_data = self.inputs.hybridization.remote_results_folder
+        transmission_data = self.inputs.transmission.remote_results_folder
+
+        if not isinstance(hybridization_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(hybridization_data)}`"
+            )
+
+        if hybridization_data.computer is None:
+            raise ValueError("Missing `Computer` node for hybridization step")
+
+        if not isinstance(transmission_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(transmission_data)}`"
+            )
+
+        if transmission_data.computer is None:
+            raise ValueError("Missing `Computer` node for transmission step")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/dft.py b/src/aiida_quantum_transport/calculations/dft.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class DFTCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.dft"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)

diff --git a/src/aiida_quantum_transport/calculations/dmft.py b/src/aiida_quantum_transport/calculations/dmft.py
@@ -2,7 +2,6 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import numpy as np
 from aiida import orm
@@ -11,17 +10,14 @@
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class DMFTCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.dmft"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -158,12 +154,8 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         matsubara_hybridization_filepath = (
             precomputed_input_dir / "matsubara_hybridization.bin"
         ).as_posix()
-        hamiltonian_filepath = (
-            precomputed_input_dir / "hamiltonian.npy"
-        ).as_posix()
-        occupancies_filepath = (
-            precomputed_input_dir / "occupancies.npy"
-        ).as_posix()
+        hamiltonian_filepath = (precomputed_input_dir / "hamiltonian.npy").as_posix()
+        occupancies_filepath = (precomputed_input_dir / "occupancies.npy").as_posix()
         scattering_region_filename = "scatt.npy"
         scattering_region_filepath = (
             precomputed_input_dir / scattering_region_filename
@@ -200,9 +192,14 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             occupancies_filepath,
         ]
 
-        hybridization_data: orm.RemoteData = (
-            self.inputs.hybridization.remote_results_folder
-        )
+        hybridization_data = self.inputs.hybridization.remote_results_folder
+
+        if not isinstance(hybridization_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(hybridization_data)}`"
+            )
+        if hybridization_data.computer is None:
+            raise ValueError("Missing `Computer` node")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/functions.py b/src/aiida_quantum_transport/calculations/functions.py
@@ -30,11 +30,7 @@ def get_scattering_region(
     y_max = y.max() if y_max is None else y_max
 
     scattering_region = np.where(
-        (z >= z_average)
-        & (x >= x_min)
-        & (x <= x_max)
-        & (y >= y_min)
-        & (y <= y_max)
+        (z >= z_average) & (x >= x_min) & (x <= x_max) & (y >= y_min) & (y <= y_max)
     )[0]
 
     return orm.ArrayData(scattering_region)
diff --git a/src/aiida_quantum_transport/calculations/greens.py b/src/aiida_quantum_transport/calculations/greens.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class GreensFunctionParametersCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.greens"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -132,12 +128,8 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
 
         precomputed_input_dir = input_dir / "precomputed"
         (temp_dir / precomputed_input_dir).mkdir()
-        leads_hamiltonian_filepath = (
-            precomputed_input_dir / "hs_leads.npy"
-        ).as_posix()
-        los_hamiltonian_filepath = (
-            precomputed_input_dir / "hs_los.npy"
-        ).as_posix()
+        leads_hamiltonian_filepath = (precomputed_input_dir / "hs_leads.npy").as_posix()
+        los_hamiltonian_filepath = (precomputed_input_dir / "hs_los.npy").as_posix()
 
         codeinfo = CodeInfo()
         codeinfo.code_uuid = self.inputs.code.uuid
@@ -156,8 +148,22 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             los_hamiltonian_filepath,
         ]
 
-        leads_data: orm.RemoteData = self.inputs.leads.remote_results_folder
-        los_data: orm.RemoteData = self.inputs.los.remote_results_folder
+        leads_data = self.inputs.leads.remote_results_folder
+        los_data = self.inputs.los.remote_results_folder
+
+        if not isinstance(leads_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(leads_data)}`"
+            )
+
+        if leads_data.computer is None:
+            raise ValueError("Missing `Computer` node for leads step")
+
+        if not isinstance(los_data, orm.RemoteData):
+            raise ValueError(f"Expected `RemoteData` instance; got `{type(los_data)}`")
+
+        if los_data.computer is None:
+            raise ValueError("Missing `Computer` node for los step")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/hybridize.py b/src/aiida_quantum_transport/calculations/hybridize.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class HybridizationCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.hybridize"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -152,9 +148,7 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
 
         precomputed_input_dir = input_dir / "precomputed"
         (temp_dir / precomputed_input_dir).mkdir()
-        los_indices_filepath = (
-            precomputed_input_dir / "los_indices.npy"
-        ).as_posix()
+        los_indices_filepath = (precomputed_input_dir / "los_indices.npy").as_posix()
         hamiltonian_ii_filepath = (
             precomputed_input_dir / "hamiltonian_ii.pkl"
         ).as_posix()
@@ -180,10 +174,22 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             self_energies_filepath,
         ]
 
-        los_data: orm.RemoteData = self.inputs.los.remote_results_folder
-        greens_function_data: orm.RemoteData = (
-            self.inputs.greens_function.remote_results_folder
-        )
+        los_data = self.inputs.los.remote_results_folder
+        greens_function_data = self.inputs.greens_function.remote_results_folder
+
+        if not isinstance(los_data, orm.RemoteData):
+            raise ValueError(f"Expected `RemoteData` instance; got `{type(los_data)}`")
+
+        if los_data.computer is None:
+            raise ValueError("Missing `Computer` node for los step")
+
+        if not isinstance(greens_function_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(greens_function_data)}`"
+            )
+
+        if greens_function_data.computer is None:
+            raise ValueError("Missing `Computer` node for greens function step")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/localize.py b/src/aiida_quantum_transport/calculations/localize.py
@@ -2,7 +2,6 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import numpy as np
 from aiida import orm
@@ -11,17 +10,14 @@
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class LocalizationCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.localize"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -125,18 +121,23 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         if self.inputs.lowdin:
             codeinfo.cmdline_params.append("--lowdin")
 
+        device_data = self.inputs.device.remote_results_folder
+
+        if not isinstance(device_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(device_data)}`"
+            )
+
+        if device_data.computer is None:
+            raise ValueError("Missing `Computer` node for leads step")
+
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]
         calcinfo.local_copy_list = []
         calcinfo.remote_symlink_list = [
             (
-                self.inputs.device.remote_results_folder.computer.uuid,
-                (
-                    Path(
-                        self.inputs.device.remote_results_folder.get_remote_path()
-                    )
-                    / restart_filename
-                ).as_posix(),
+                device_data.computer.uuid,
+                f"{device_data.get_remote_path()}/{restart_filename}",
                 restart_filepath,
             )
         ]