From 33458f322427cf4aa58ac642771942f19eb7970c Mon Sep 17 00:00:00 2001 From: Edan Bainglass Date: Mon, 5 Feb 2024 11:39:09 +0000 Subject: [PATCH] Implement transmission step --- pyproject.toml | 2 + .../calculations/__init__.py | 2 + .../calculations/transmission.py | 183 ++++++++++++++++++ .../parsers/__init__.py | 2 + .../parsers/transmission.py | 19 ++ .../workchains/coulomb_diamonds.py | 51 ++++- 6 files changed, 258 insertions(+), 1 deletion(-) create mode 100644 src/aiida_quantum_transport/calculations/transmission.py create mode 100644 src/aiida_quantum_transport/parsers/transmission.py diff --git a/pyproject.toml b/pyproject.toml index ace28f1..7fcde5c 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/aiida_quantum_transport/calculations/__init__.py b/src/aiida_quantum_transport/calculations/__init__.py index 9f4b235..aae11d2 100644 --- a/src/aiida_quantum_transport/calculations/__init__.py +++ b/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", ] diff --git a/src/aiida_quantum_transport/calculations/transmission.py b/src/aiida_quantum_transport/calculations/transmission.py new file mode 100644 index 0000000..0db52dd --- /dev/null +++ b/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 diff --git a/src/aiida_quantum_transport/parsers/__init__.py b/src/aiida_quantum_transport/parsers/__init__.py index 278d4f3..84b39e8 100644 --- a/src/aiida_quantum_transport/parsers/__init__.py +++ b/src/aiida_quantum_transport/parsers/__init__.py @@ -3,6 +3,7 @@ from .dmft import DMFTParser from .hybridize import HybridizationParser from .localize import LocalizationParser +from .transmission import TransmissionParser __all__ = [ "CustomParser", @@ -10,4 +11,5 @@ "LocalizationParser", "HybridizationParser", "DMFTParser", + "TransmissionParser", ] diff --git a/src/aiida_quantum_transport/parsers/transmission.py b/src/aiida_quantum_transport/parsers/transmission.py new file mode 100644 index 0000000..cd0064a --- /dev/null +++ b/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 diff --git a/src/aiida_quantum_transport/workchains/coulomb_diamonds.py b/src/aiida_quantum_transport/workchains/coulomb_diamonds.py index cc2c5ca..5e90526 100644 --- a/src/aiida_quantum_transport/workchains/coulomb_diamonds.py +++ b/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,6 +140,11 @@ 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, @@ -140,7 +152,7 @@ def define(cls, spec: WorkChainSpec) -> None: 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", + ) + )