From 021cee6ee4a10e76b82044d677ec27f28ca14a4e Mon Sep 17 00:00:00 2001 From: Edan Bainglass Date: Sat, 27 Jan 2024 16:30:03 +0000 Subject: [PATCH] Implement hybridization step --- pyproject.toml | 2 + .../calculations/__init__.py | 2 + .../calculations/hybridize.py | 208 ++++++++++++++++++ .../calculations/localize.py | 3 +- .../parsers/__init__.py | 2 + .../parsers/hybridize.py | 32 +++ .../workchains/coulomb_diamonds.py | 51 ++++- 7 files changed, 295 insertions(+), 5 deletions(-) create mode 100644 src/aiida_quantum_transport/calculations/hybridize.py create mode 100644 src/aiida_quantum_transport/parsers/hybridize.py diff --git a/pyproject.toml b/pyproject.toml index aa70438..85dcffb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ requires-python = ">=3.9" "quantum_transport" = "aiida_quantum_transport.calculations.custom:CustomCalculation" "quantum_transport.dft" = "aiida_quantum_transport.calculations.dft:DftCalculation" "quantum_transport.localize" = "aiida_quantum_transport.calculations.localize:LocalizationCalculation" +"quantum_transport.hybridize" = "aiida_quantum_transport.calculations.hybridize:HybridizationCalculation" [project.entry-points."aiida.cmdline.data"] "quantum_transport" = "aiida_quantum_transport.cli.commands:data_cli" @@ -42,6 +43,7 @@ requires-python = ">=3.9" "quantum_transport" = "aiida_quantum_transport.parsers.custom:CustomParser" "quantum_transport.dft" = "aiida_quantum_transport.parsers.dft:DftParser" "quantum_transport.localize" = "aiida_quantum_transport.parsers.localize:LocalizationParser" +"quantum_transport.hybridize" = "aiida_quantum_transport.parsers.hybridize:HybridizationParser" [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 582705d..7397531 100644 --- a/src/aiida_quantum_transport/calculations/__init__.py +++ b/src/aiida_quantum_transport/calculations/__init__.py @@ -1,6 +1,7 @@ from .custom import CustomCalculation from .dft import DftCalculation from .functions import * +from .hybridize import HybridizationCalculation from .localize import LocalizationCalculation __all__ = [ @@ -8,4 +9,5 @@ "DftCalculation", "get_scattering_region", "LocalizationCalculation", + "HybridizationCalculation", ] diff --git a/src/aiida_quantum_transport/calculations/hybridize.py b/src/aiida_quantum_transport/calculations/hybridize.py new file mode 100644 index 0000000..adb25c0 --- /dev/null +++ b/src/aiida_quantum_transport/calculations/hybridize.py @@ -0,0 +1,208 @@ +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 HybridizationCalculation(CalcJob): + """docstring""" + + _default_parser_name = "quantum_transport.hybridize" + + @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( + "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( + "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( + "metadata.options.parser_name", + valid_type=str, + default=cls._default_parser_name, + ) + + spec.output( + "hybridization_file", + valid_type=orm.SinglefileData, + help="", # TODO fill in + ) + + spec.output( + "energies_file", + valid_type=orm.SinglefileData, + help="The energies file", + ) + + spec.output( + "hamiltonian_file", + valid_type=orm.SinglefileData, + help="The hamiltonian file", + ) + + spec.output( + "eigenvalues_file", + valid_type=orm.SinglefileData, + help="The eigenvalues file", + ) + + spec.output( + "matsubara_hybridization_file", + valid_type=orm.SinglefileData, + help="The Matsubara hybridization file", + ) + + spec.output( + "matsubara_energies_file", + valid_type=orm.SinglefileData, + help="The Matsubara energies file", + ) + + spec.output( + "occupancies_file", + valid_type=orm.SinglefileData, + help="The occupancies file", + ) + + 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) + + 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, + ] + + 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, + ), + ] + calcinfo.retrieve_list = [ + "hybridization.bin", + "energies.npy", + "hamiltonian.npy", + "eigenvalues.npy", + "matsubara_hybridization.bin", + "matsubara_energies.npy", + "occupancies.npy", + ] + + return calcinfo diff --git a/src/aiida_quantum_transport/calculations/localize.py b/src/aiida_quantum_transport/calculations/localize.py index d05d7e1..8c33380 100644 --- a/src/aiida_quantum_transport/calculations/localize.py +++ b/src/aiida_quantum_transport/calculations/localize.py @@ -40,6 +40,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None: "scattering.region", valid_type=orm.ArrayData, required=False, + default=lambda: orm.ArrayData([]), help="The scattering region", ) @@ -52,7 +53,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None: spec.input( "lowdin", valid_type=orm.Bool, - default=orm.Bool(False), + default=lambda: orm.Bool(False), help="", # TODO fill in ) diff --git a/src/aiida_quantum_transport/parsers/__init__.py b/src/aiida_quantum_transport/parsers/__init__.py index 13bf94f..2027ddd 100644 --- a/src/aiida_quantum_transport/parsers/__init__.py +++ b/src/aiida_quantum_transport/parsers/__init__.py @@ -1,9 +1,11 @@ from .custom import CustomParser from .dft import DftParser +from .hybridize import HybridizationParser from .localize import LocalizationParser __all__ = [ "CustomParser", "DftParser", "LocalizationParser", + "HybridizationParser", ] diff --git a/src/aiida_quantum_transport/parsers/hybridize.py b/src/aiida_quantum_transport/parsers/hybridize.py new file mode 100644 index 0000000..92bf501 --- /dev/null +++ b/src/aiida_quantum_transport/parsers/hybridize.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path + +from aiida import orm +from aiida.engine import ExitCode +from aiida.parsers import Parser + + +class HybridizationParser(Parser): + """docstring""" + + _OUTPUT_FILE_LIST = [ + "hybridization.bin", + "energies.npy", + "hamiltonian.npy", + "eigenvalues.npy", + "matsubara_hybridization.bin", + "matsubara_energies.npy", + "occupancies.npy", + ] + + def parse(self, **kwargs) -> ExitCode | None: + """docstring""" + + for filename in self._OUTPUT_FILE_LIST: + path = Path(self.node.get_remote_workdir()) / filename + prefix = filename.split(".")[0] + output_label = f"{prefix}_file" + self.out(output_label, orm.SinglefileData(path)) + + return None diff --git a/src/aiida_quantum_transport/workchains/coulomb_diamonds.py b/src/aiida_quantum_transport/workchains/coulomb_diamonds.py index ccd89c7..7c6a3b2 100644 --- a/src/aiida_quantum_transport/workchains/coulomb_diamonds.py +++ b/src/aiida_quantum_transport/workchains/coulomb_diamonds.py @@ -7,6 +7,7 @@ from aiida_quantum_transport.calculations import ( DftCalculation, + HybridizationCalculation, LocalizationCalculation, get_scattering_region, ) @@ -48,11 +49,12 @@ def define(cls, spec: WorkChainSpec) -> None: exclude=["code"], ) + # TODO rethink this one (redefines localization input) spec.input( "scattering.region", valid_type=orm.Dict, required=False, - default=orm.Dict({}), + default=lambda: orm.Dict({}), help="The xy-limits defining the scattering region", ) @@ -68,6 +70,12 @@ def define(cls, spec: WorkChainSpec) -> None: include=["code", "lowdin", "metadata"], ) + spec.expose_inputs( + HybridizationCalculation, + namespace="hybridization", + include=["code", "basis", "parameters", "metadata"], + ) + spec.expose_outputs( DftCalculation, namespace="dft.leads", @@ -83,11 +91,16 @@ def define(cls, spec: WorkChainSpec) -> None: namespace="localization", ) + spec.expose_outputs( + HybridizationCalculation, + namespace="hybridization", + ) + spec.outline( cls.run_dft, cls.define_scattering_region, cls.transform_basis, - # cls.compute_hybridization, + cls.compute_hybridization, # cls.run_dmft_adjust_mu, # cls.run_dmft_sweep_mu, # cls.compute_transmission, @@ -111,8 +124,6 @@ def run_dft(self): return ToContext( dft_leads=self.submit(DftCalculation, **leads_inputs), dft_device=self.submit(DftCalculation, **device_inputs), - # dft_leads=orm.load_node(1703), - # dft_device=orm.load_node(1704), ) def define_scattering_region(self): @@ -144,6 +155,30 @@ def transform_basis(self): def compute_hybridization(self): """docstring""" + hybridization_inputs = { + "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, + }, + **self.exposed_inputs( + HybridizationCalculation, + namespace="hybridization", + ), + } + return ToContext( + hybridization=self.submit( + HybridizationCalculation, + **hybridization_inputs, + ) + ) def run_dmft_adjust_mu(self): """docstring""" @@ -183,3 +218,11 @@ def gather_results(self): namespace="localization", ) ) + + self.out_many( + self.exposed_outputs( + self.ctx.hybridization, + HybridizationCalculation, + namespace="hybridization", + ) + )