Pyiron job for Extra FIM simulations

.gitignore

@@ -2,3 +2,6 @@
 *.swo
 __pycache__/
 .idea/
+*.lock
+.vscode/
+

data_utils/plotting.py → EXTRA_FIM/datautils/plotting.py

@@ -159,7 +159,7 @@ def waves_figure(
     # --- compute and plot extrapolated wave
     if compute_extra:
         print("Computing EXTRA wave...")
-        from ..EXTRA_FIM.extra import extra_waves
+        from EXTRA_FIM.extra import extra_waves
 
         waves_extra = extra_waves(Simulator, reader=waves_reader, pot=pot)
         _, psi_extra = waves_extra.get_psi(istate, ispin, ik)

data_utils/pre_processing.py → EXTRA_FIM/datautils/pre_processing.py

@@ -13,7 +13,7 @@
 import scipy
 
 from .plotting import potential_figure
-from ..EXTRA_FIM.potential import sx_el_potential1D_cell
+from EXTRA_FIM.potential import sx_el_potential1D_cell
 
 
 class PreProcessingFIM:

EXTRA_FIM/extra.py

@@ -1,11 +1,9 @@
+from pathlib import Path
 import numpy as np
-import netCDF4
 import scipy.constants
 import scipy.optimize
-from .waves_reader_abc import waves_reader_abc
-from .sx_nc_waves_reader import sx_nc_waves_reader
-from pathlib import Path
-
+from EXTRA_FIM.waves_reader_abc import waves_reader_abc
+from EXTRA_FIM.sx_nc_waves_reader import sx_nc_waves_reader
 
 __author__ = "Shalini Bhatt"
 __copyright__ = (

EXTRA_FIM/main.py

@@ -1,10 +1,10 @@
+from pathlib import Path
 import numpy as np
 import h5py
 import netCDF4
 import scipy.constants
 import scipy.optimize
-from pathlib import Path
-from .extra import extra_waves
+from EXTRA_FIM.extra import extra_waves
 
 
 __author__ = "Shalini Bhatt"
@@ -116,7 +116,7 @@ def sum_all_states(self):
         for ik in range(0, self.wf.nk):
             sum_single_k(ik)
 
-    def sum_single_k(self, ik):
+    def sum_single_k(self, ik, **kwargs):
         """compute partial fim image for several ionization energies for all eigenstates between Efermi
         and Emax for one k point (ik). Save the partial dos files"""
 
@@ -175,7 +175,10 @@ def sum_single_k(self, ik):
                     )
 
         # --- write output file
-        filename = f"partial_dos{ik}.h5"
+        if 'path' in kwargs:
+            filename = kwargs['path']+'/'+f"partial_dos{ik}.h5"
+        else:
+            filename = f"partial_dos{ik}.h5"
         with h5py.File(filename, "w") as handle:
             handle.create_dataset(
                 "ionization_energies", data=self.inputDict["ionization_energies"]

EXTRA_FIM/pyiron_job.py

@@ -0,0 +1,247 @@
+from pathlib import Path
+from pyiron_base.utils.error import ImportAlarm
+from pyiron_base.jobs.master.parallel import ParallelMaster
+from pyiron_base.jobs.job.jobtype import JobType
+from pyiron_base import JobGenerator
+from pyiron_base.jobs.job.template import TemplateJob
+import h5py
+import numpy as np
+from pyiron_atomistics.sphinx.structure import read_atoms
+
+try:
+    import EXTRA_FIM.main as fim
+    from EXTRA_FIM.potential import extend_potential, sx_el_potential3D_cell
+    from EXTRA_FIM.datautils.pre_processing import suggest_input_dictionary
+    from EXTRA_FIM.sx_nc_waves_reader import sx_nc_waves_reader
+except ImportError:
+    import_alarm = ImportAlarm("Unable to import EXTRA_FIM")
+
+__author__ = "Christoph Freysoldt, Shyam Katnagallu"
+__copyright__ = (
+    "Copyright 2021, Max-Planck-Institut für Eisenforschung GmbH - "
+    "Computational Materials Design (CM) Department"
+)
+__version__ = "0.1.0"
+__maintainer__ = "Shyam Katnagallu"
+__email__ = "[email protected]"
+__status__ = "development"
+__date__ = "March 5, 2024"
+
+
+class ExtraFimSimulatorRefJob(TemplateJob):
+    """Reference pyiron job for Extra FIM simulation."""
+
+    def __init__(self, project, job_name):
+        super().__init__(project, job_name)
+        self.input["waves_directory"] = None
+        self.input["waves_reader"] = None
+        self.input["kpoint"] = None
+        self.input["ionization_energies"] = None
+        self.input["extrapolate_potential"] = False
+        self.input["total_kpoints"] = None
+
+    def extrpolate_potential(self):
+        """Extrapolate potential if needed, to extrapolate waves to higher distances"""
+
+        elec_potential, _ = sx_el_potential3D_cell(
+            self.input["simulator_dict"]["working_directory"]
+        )
+        pot, _, _, cell = fim.potential(self.input["simulator_dict"]).potential_cell()
+
+        if self.input.extrapolate_potential:
+            iz0 = self.input["simulator_dict"]["iz_ext_from"]
+            new_z_max = self.input["simulator_dict"]["z_ext"]
+            _, pot_ext = extend_potential(
+                elec_potential / fim.HARTREE_TO_EV,
+                iz0,
+                pot,
+                cell,
+                z_max=new_z_max,
+                izend=self.input["simulator_dict"]["izend"],
+                dv_limit=1e-4,
+                plotG=1,
+            )
+            # copy extension from pot to elec_potential
+            elec_ext = pot_ext[:, :, :, 0] * fim.HARTREE_TO_EV
+            elec_ext[:, :, 0:iz0] = elec_potential[:, :, :iz0]
+        else:
+            pot_ext = pot
+            elec_ext = elec_potential
+        return pot_ext, elec_potential
+
+    @property
+    def suggest_input_dict(self):
+        """Suggests a input dictionary based on the electrostatic potential, Fermi and ionization energy"""
+        waves_reader = sx_nc_waves_reader(
+            Path(self.input["waves_directory"]) / "waves.sxb"
+        )
+        e_fermi = waves_reader.get_fermi_energy()
+        _, sim = suggest_input_dictionary(
+            self.input["waves_directory"],
+            e_fermi,
+            ionization_energies=self.input["ionization_energies"],
+        )
+        self.input["total_kpoints"] = waves_reader.nk
+        self.input["simulator_dict"] = sim
+        self.input["z_max"] = sim["z_max"]  # rename later
+        self.input["izstart_min"] = sim["izstart_min"]  # rename later
+        self.input["izend"] = sim["izend"]  # rename later
+        self.input["limit"] = sim["limit"]  # rename later
+        self.input["cutoff"] = sim["cutoff"]  # rename later
+        self.input["E_fermi"] = sim["E_fermi"]  # rename later
+        self.input["E_max"] = sim["E_max"]  # rename later
+        return sim
+
+    def run_static(self):
+
+        self.project_hdf5.create_working_directory()
+        # self.suggest_input_dict
+        pot_ext, elec_ext = self.extrpolate_potential()
+        waves_reader = sx_nc_waves_reader(self.input["waves_directory"] + "/waves.sxb")
+        fimsim = fim.FIM_simulations(
+            self.input["simulator_dict"],
+            reader=waves_reader,
+            V_total=pot_ext,
+            V_elstat=elec_ext,
+        )
+        self.status.submitted = True
+        fimsim.sum_single_k(self.input["kpoint"], path=self.working_directory)
+        self.status.finished = True
+
+
+class ExtraFimSimulatorJobGenerator(JobGenerator):
+    """Job generator class for extra fim simulator pyiron jobs"""
+
+    @property
+    def parameter_list(self):
+        """
+
+        Returns:
+            (list)
+        """
+        parameter_lst = []
+        kpoints = self._master.input.get("kpoint")
+        if kpoints is None:
+            waves_reader = sx_nc_waves_reader(
+                Path(self._master.input["waves_directory"]) / "waves.sxb"
+            )
+        for k in range(waves_reader.nk):
+            parameter_lst.append(k)
+        return parameter_lst
+
+    def job_name(self, parameter):
+        k_point = parameter
+        return f"{self._master.job_name}_kpoint_{k_point}"
+
+    def modify_job(self, job, parameter):
+        job.structure = self._master.get_structure
+        for k in self._master.input.keys():
+            job.input[k] = self._master.input[k]
+        job.input.kpoint = parameter
+        return job
+
+
+class ExtraFimSimulator(ParallelMaster):
+    """ "Pyiron Extra FIM simulator job class to make subjobs for each k point"""
+
+    def __init__(self, project, job_name):
+        super(ExtraFimSimulator, self).__init__(project, job_name)
+        self.__version__ = "0.1.0"
+        self.input["waves_directory"] = None
+        self.input["waves_reader"] = None
+        self.input["kpoint"] = None
+        self.input["ionization_energies"] = None
+        self.input["extrapolate_potential"] = False
+        self.input["total_kpoints"] = None
+        self._job_generator = ExtraFimSimulatorJobGenerator(self)
+        self.ref_job = ExtraFimSimulatorRefJob(project=project, job_name=job_name)
+
+    def extrpolate_potential(self):
+        """Extrapolate potential if needed, to extrapolate waves to higher distances"""
+
+        elec_potential, _ = sx_el_potential3D_cell(
+            self.input["simulator_dict"]["working_directory"]
+        )
+        pot, _, _, cell = fim.potential(self.input["simulator_dict"]).potential_cell()
+
+        if self.input.extrapolate_potential:
+            iz0 = self.input["simulator_dict"]["iz_ext_from"]
+            new_z_max = self.input["simulator_dict"]["z_ext"]
+            _, pot_ext = extend_potential(
+                elec_potential / fim.HARTREE_TO_EV,
+                iz0,
+                pot,
+                cell,
+                z_max=new_z_max,
+                izend=self.input["simulator_dict"]["izend"],
+                dv_limit=1e-4,
+                plotG=1,
+            )
+            # copy extension from pot to elec_potential
+            elec_ext = pot_ext[:, :, :, 0] * fim.HARTREE_TO_EV
+            elec_ext[:, :, 0:iz0] = elec_potential[:, :, :iz0]
+        else:
+            pot_ext = pot
+            elec_ext = elec_potential
+        return pot_ext, elec_potential
+
+    @property
+    def suggest_input_dict(self):
+        """Suggests a input dictionary based on the electrostatic potential, Fermi and ionization energy and populates input"""
+        waves_reader = sx_nc_waves_reader(
+            Path(self.input["waves_directory"]) / "waves.sxb"
+        )
+        e_fermi = waves_reader.get_fermi_energy()
+        _, sim = suggest_input_dictionary(
+            self.input["waves_directory"],
+            e_fermi,
+            ionization_energies=self.input["ionization_energies"],
+        )
+        self.input["total_kpoints"] = waves_reader.nk
+        self.input["simulator_dict"] = sim
+        self.input["z_max"] = sim["z_max"]  # rename later
+        self.input["izstart_min"] = sim["izstart_min"]  # rename later
+        self.input["izend"] = sim["izend"]  # rename later
+        self.input["limit"] = sim["limit"]  # rename later
+        self.input["cutoff"] = sim["cutoff"]  # rename later
+        self.input["E_fermi"] = sim["E_fermi"]  # rename later
+        self.input["E_max"] = sim["E_max"]  # rename later
+        self.get_structure
+        return sim
+
+    @property
+    def get_structure(self):
+        """Tries to get the relaxed sphinx structure if available"""
+        if (Path(self.input["waves_directory"]) / "relaxedStr.sx").exists():
+            self.structure = read_atoms(
+                Path(self.input["waves_directory"]) / "relaxedStr.sx"
+            )
+
+    def collect_output(self):
+        FIM_total = {}
+        zFIM_total = {}
+        for job_id in self.child_ids:
+            subjob = self.project_hdf5.load(job_id)
+            subjobdir = subjob.working_directory
+            ik = subjob.input.kpoint
+            IEs = subjob.input.ionization_energies
+            with h5py.File(f"{subjobdir}/partial_dos{ik}.h5") as handle:
+                for IE in IEs:
+                    fim_k = np.asarray(handle[f"IE={IE}"])
+                    zfim_k = np.asarray(handle[f"zIE={IE}"])
+                    if IE not in FIM_total:
+                        FIM_total[IE] = np.zeros_like(fim_k)
+                        zFIM_total[IE] = np.zeros_like(zfim_k)
+                    FIM_total[IE] += fim_k
+                    zFIM_total[IE] += zfim_k
+
+        for IE in IEs:
+            self._output[f"total_FIM/{IE}"] = FIM_total[IE]
+            self._output[f"z_resolved_FIM/{IE}"] = zFIM_total[IE]
+
+        with self.project_hdf5.open("output") as hdf5_out:
+            for key, val in self._output.items():
+                hdf5_out[key] = val
+
+
+JobType.register(ExtraFimSimulator)

Examples/FIMjob.ipynb

@@ -16,13 +16,11 @@
    "outputs": [],
    "source": [
     "from pyiron_base import load\n",
-    "import numpy as np\n",
-    "\n",
     "import sys\n",
     "sys.path.insert(0,'/cmmc/u/cfrey/devel/python-test')\n",
-    "import EXTRA_FIM.EXTRA_FIM.main as fim\n",
-    "from EXTRA_FIM.EXTRA_FIM.sx_nc_waves_reader import sx_nc_waves_reader\n",
-    "from EXTRA_FIM.EXTRA_FIM.potential import extend_potential, sx_el_potential3D_cell"
+    "import EXTRA_FIM.main as fim\n",
+    "from EXTRA_FIM.sx_nc_waves_reader import sx_nc_waves_reader\n",
+    "from EXTRA_FIM.potential import extend_potential, sx_el_potential3D_cell"
    ]
   },
   {
@@ -157,7 +155,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.8"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,