Implement transmission step

pyproject.toml

@@ -33,6 +33,7 @@ requires-python = ">=3.9"
 "quantum_transport.localize" = "aiida_quantum_transport.calculations.localize:LocalizationCalculation"
 "quantum_transport.hybridize" = "aiida_quantum_transport.calculations.hybridize:HybridizationCalculation"
 "quantum_transport.dmft" = "aiida_quantum_transport.calculations.dmft:DMFTCalculation"
+"quantum_transport.transmission" = "aiida_quantum_transport.calculations.transmission:TransmissionCalculation"
 
 [project.entry-points."aiida.cmdline.data"]
 "quantum_transport" = "aiida_quantum_transport.cli.commands:data_cli"
@@ -46,6 +47,7 @@ requires-python = ">=3.9"
 "quantum_transport.localize" = "aiida_quantum_transport.parsers.localize:LocalizationParser"
 "quantum_transport.hybridize" = "aiida_quantum_transport.parsers.hybridize:HybridizationParser"
 "quantum_transport.dmft" = "aiida_quantum_transport.parsers.dmft:DMFTParser"
+"quantum_transport.transmission" = "aiida_quantum_transport.parsers.transmission:TransmissionParser"
 
 [project.optional-dependencies]
 docs = ["sphinx-design~=0.4.1", "pydata-sphinx-theme==0.13.3"]

src/aiida_quantum_transport/calculations/__init__.py

@@ -4,6 +4,7 @@
 from .functions import *
 from .hybridize import HybridizationCalculation
 from .localize import LocalizationCalculation
+from .transmission import TransmissionCalculation
 
 __all__ = [
     "CustomCalculation",
@@ -12,4 +13,5 @@
     "LocalizationCalculation",
     "HybridizationCalculation",
     "DMFTCalculation",
+    "TransmissionCalculation",
 ]

src/aiida_quantum_transport/calculations/transmission.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import pickle
+from typing import TYPE_CHECKING
+
+from aiida import orm
+from aiida.common.datastructures import CalcInfo, CodeInfo
+from aiida.common.folders import Folder
+from aiida.engine import CalcJob
+
+if TYPE_CHECKING:
+    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
+
+
+class TransmissionCalculation(CalcJob):
+    """docstring"""
+
+    _default_parser_name = "quantum_transport.transmission"
+
+    @classmethod
+    def define(cls, spec: CalcJobProcessSpec) -> None:
+        """docstring"""
+
+        super().define(spec)
+
+        spec.input(
+            "code",
+            valid_type=orm.AbstractCode,
+            help="The hybridization script",
+        )
+
+        spec.input(
+            "leads.structure",
+            valid_type=orm.StructureData,
+            help="The structure of the leads",
+        )
+
+        spec.input(
+            "leads.kpoints",
+            valid_type=orm.KpointsData,
+            help="The kpoints mesh used for the leads",
+        )
+
+        spec.input(
+            "leads.hamiltonian_file",
+            valid_type=orm.SinglefileData,
+            help="The file holding the leads hamiltonian",
+        )
+
+        spec.input(
+            "device.structure",
+            valid_type=orm.StructureData,
+            help="The structure of the device",
+        )
+
+        spec.input(
+            "basis",
+            valid_type=orm.Dict,
+            help="",  # TODO fill in
+        )
+
+        spec.input(
+            "parameters",
+            valid_type=orm.Dict,
+            required=False,
+            default=lambda: orm.Dict({}),
+            help="parameters used for orbital hybridization",
+        )
+
+        spec.input(
+            "localization.hamiltonian_file",
+            valid_type=orm.SinglefileData,
+            help="The file holding the localized scattering hamiltonian",
+        )
+
+        spec.input(
+            "localization.index_file",
+            valid_type=orm.SinglefileData,
+            help="",  # TODO fill in
+        )
+
+        spec.input(
+            "dmft.sigma_folder",
+            valid_type=orm.FolderData,
+            help="folder containing self-energy files",
+        )
+
+        spec.input(
+            "metadata.options.parser_name",
+            valid_type=str,
+            default=cls._default_parser_name,
+        )
+
+        spec.output(
+            "transmission_folder",
+            valid_type=orm.FolderData,
+            help="The transmission folder",
+        )
+
+    def prepare_for_submission(self, folder: Folder) -> CalcInfo:
+        """docstring"""
+
+        pickled_leads_structure_filename = "leads_structure.pkl"
+        with folder.open(pickled_leads_structure_filename, "wb") as file:
+            leads: orm.StructureData = self.inputs.leads.structure
+            pickle.dump(leads.get_ase(), file)
+
+        pickled_device_structure_filename = "device_structure.pkl"
+        with folder.open(pickled_device_structure_filename, "wb") as file:
+            device: orm.StructureData = self.inputs.device.structure
+            pickle.dump(device.get_ase(), file)
+
+        pickled_leads_kpoints_filename = "leads_kpoints.pkl"
+        with folder.open(pickled_leads_kpoints_filename, "wb") as file:
+            kpoints: orm.KpointsData = self.inputs.leads.kpoints
+            pickle.dump(kpoints.get_kpoints_mesh()[0], file)
+
+        pickled_basis_filename = "basis.pkl"
+        with folder.open(pickled_basis_filename, "wb") as file:
+            basis: orm.Dict = self.inputs.basis
+            pickle.dump(basis.get_dict(), file)
+
+        pickled_parameters_filename = "parameters.pkl"
+        with folder.open(pickled_parameters_filename, "wb") as file:
+            parameters: orm.Dict = self.inputs.parameters
+            pickle.dump(parameters.get_dict(), file)
+
+        # HANDLE sigma folder
+
+        codeinfo = CodeInfo()
+        codeinfo.code_uuid = self.inputs.code.uuid
+        codeinfo.cmdline_params = [
+            "--leads-structure-filename",
+            pickled_leads_structure_filename,
+            "--device-structure-filename",
+            pickled_device_structure_filename,
+            "--leads-kpoints-filename",
+            pickled_leads_kpoints_filename,
+            "--leads-hamiltonian-filename",
+            f"leads_{self.inputs.leads.hamiltonian_file.filename}",
+            "--localized-hamiltonian-filename",
+            f"device_{self.inputs.localization.hamiltonian_file.filename}",
+            "--localization-index-filename",
+            self.inputs.localization.index_file.filename,
+            "--basis-filename",
+            pickled_basis_filename,
+            "--parameters-filename",
+            pickled_parameters_filename,
+            "--sigma-folder-name",
+            "sigma_folder",
+        ]
+
+        self.node.get_remote_workdir()
+
+        calcinfo = CalcInfo()
+        calcinfo.codes_info = [codeinfo]
+        calcinfo.local_copy_list = [
+            (
+                self.inputs.leads.hamiltonian_file.uuid,
+                self.inputs.leads.hamiltonian_file.filename,
+                f"leads_{self.inputs.leads.hamiltonian_file.filename}",
+            ),
+            (
+                self.inputs.localization.hamiltonian_file.uuid,
+                self.inputs.localization.hamiltonian_file.filename,
+                f"device_{self.inputs.localization.hamiltonian_file.filename}",
+            ),
+            (
+                self.inputs.localization.index_file.uuid,
+                self.inputs.localization.index_file.filename,
+                self.inputs.localization.index_file.filename,
+            ),
+            (
+                self.inputs.dmft.sigma_folder.uuid,
+                ".",
+                "sigma_folder",
+            ),
+        ]
+        calcinfo.retrieve_list = [
+            "transmission_folder",
+        ]
+
+        return calcinfo

src/aiida_quantum_transport/parsers/__init__.py

@@ -3,11 +3,13 @@
 from .dmft import DMFTParser
 from .hybridize import HybridizationParser
 from .localize import LocalizationParser
+from .transmission import TransmissionParser
 
 __all__ = [
     "CustomParser",
     "DFTParser",
     "LocalizationParser",
     "HybridizationParser",
     "DMFTParser",
+    "TransmissionParser",
 ]

src/aiida_quantum_transport/parsers/transmission.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from aiida import orm
+from aiida.engine import ExitCode
+from aiida.parsers import Parser
+
+
+class TransmissionParser(Parser):
+    """docstring"""
+
+    def parse(self, **kwargs) -> ExitCode | None:
+        """docstring"""
+
+        path = Path(self.node.get_remote_workdir()) / "transmission_folder"
+        self.out("transmission_folder", orm.FolderData(tree=path))
+
+        return None

src/aiida_quantum_transport/workchains/coulomb_diamonds.py

@@ -10,6 +10,7 @@
     DMFTCalculation,
     HybridizationCalculation,
     LocalizationCalculation,
+    TransmissionCalculation,
     get_scattering_region,
 )
 
@@ -103,6 +104,12 @@ def define(cls, spec: WorkChainSpec) -> None:
             help="The chemical potential sweep parameters",
         )
 
+        spec.expose_inputs(
+            TransmissionCalculation,
+            namespace="transmission",
+            include=["code", "metadata"],
+        )
+
         spec.expose_outputs(
             DFTCalculation,
             namespace="dft.leads",
@@ -133,14 +140,19 @@ def define(cls, spec: WorkChainSpec) -> None:
             namespace="dmft.sweep_mu",
         )
 
+        spec.expose_outputs(
+            TransmissionCalculation,
+            namespace="transmission",
+        )
+
         spec.outline(
             cls.run_dft,
             cls.define_scattering_region,
             cls.transform_basis,
             cls.compute_hybridization,
             cls.run_dmft_converge_mu,
             cls.run_dmft_sweep_mu,
-            # cls.compute_transmission,
+            cls.compute_transmission,
             # cls.compute_current,
             cls.gather_results,
         )
@@ -289,6 +301,35 @@ def run_dmft_sweep_mu(self):
 
     def compute_transmission(self):
         """docstring"""
+        transmission_inputs = {
+            "basis": self.ctx.hybridization.inputs.basis,
+            "parameters": self.ctx.hybridization.inputs.parameters,
+            "leads": {
+                "structure": self.inputs.dft.leads.structure,
+                "kpoints": self.inputs.dft.leads.kpoints,
+                "hamiltonian_file": self.ctx.dft_leads.outputs.hamiltonian_file,
+            },
+            "device": {
+                "structure": self.inputs.dft.device.structure,
+            },
+            "localization": {
+                "index_file": self.ctx.localization.outputs.index_file,
+                "hamiltonian_file": self.ctx.localization.outputs.hamiltonian_file,
+            },
+            "dmft": {
+                "sigma_folder": self.ctx.dmft_sweep_mu.outputs.sigma_folder,
+            },
+            **self.exposed_inputs(
+                TransmissionCalculation,
+                namespace="transmission",
+            ),
+        }
+        return ToContext(
+            transmission=self.submit(
+                TransmissionCalculation,
+                **transmission_inputs,
+            )
+        )
 
     def compute_current(self):
         """docstring"""
@@ -343,3 +384,11 @@ def gather_results(self):
                 namespace="dmft.sweep_mu",
             )
         )
+
+        self.out_many(
+            self.exposed_outputs(
+                self.ctx.transmission,
+                TransmissionCalculation,
+                namespace="transmission",
+            )
+        )