Support for 1D and 2D systems with kcp.x

CITATION.cff

@@ -1,6 +1,6 @@
 authors:
-- affiliation: École Polytechnique Fédérale de Lausanne
-  email: edward.linscott@epfl.ch
+- affiliation: Paul Scherrer Institute
+  email: edward.linscott@psi.ch
   family-names: Linscott
   given-names: Edward
   orcid: https://orcid.org/0000-0002-4967-9873

pyproject.toml

@@ -5,7 +5,7 @@ description = "Koopmans spectral functional calculations with python and Quantum
 readme = "README_pypi.rst"
 requires-python = ">=3.8"
 authors = [
-    { name = "Edward Linscott", email = "edward.linscott@epfl.ch" },
+    { name = "Edward Linscott", email = "edward.linscott@psi.ch" },
     { name = "Riccardo De Gennaro", email = "riccardo.degennaro@epfl.ch" },
     { name = "Nicola Colonna", email = "nicola.colonna@psi.ch" },
 ]

src/koopmans/calculators/_projwfc.py

@@ -75,7 +75,8 @@ def generate_dos(self) -> GridDOSCollection:
             # The filename does not encode the principal quantum number n. In order to recover this number, we compare
             # the reported angular momentum quantum number l against the list of expected orbitals, and infer n
             # assuming only that the file corresponding to nl will come before (n+1)l
-            orbitals = copy.copy(self._expected_orbitals[atom.symbol])
+            label = atom.symbol + str(atom.tag) if atom.tag > 0 else atom.symbol
+            orbitals = copy.copy(self._expected_orbitals[label])
             for filename in filenames:
                 # Infer l from the filename
                 subshell = filename.name[-2]

src/koopmans/pseudopotentials.py

@@ -160,14 +160,15 @@ def expected_subshells(atoms: Atoms, pseudopotentials: Dict[str, str],
 
     expected_orbitals = {}
     for atom in atoms:
-        if atom.symbol in expected_orbitals:
+        label = atom.symbol + str(atom.tag) if atom.tag > 0 else atom.symbol
+        if label in expected_orbitals:
             continue
-        pseudo_file = pseudopotentials[atom.symbol]
+        pseudo_file = pseudopotentials[label]
         z_core = atom.number - valence_from_pseudo(pseudo_file, pseudo_dir)
         if z_core in z_core_to_first_orbital:
             first_orbital = z_core_to_first_orbital[z_core]
         else:
             raise ValueError(f'Failed to identify the subshells of the valence of {pseudo_file}')
         all_orbitals = list(z_core_to_first_orbital.values()) + ['5d', '6p', '6d']
-        expected_orbitals[atom.symbol] = sorted(all_orbitals[all_orbitals.index(first_orbital):])
+        expected_orbitals[label] = sorted(all_orbitals[all_orbitals.index(first_orbital):])
     return expected_orbitals

src/koopmans/utils/_io.py

@@ -52,7 +52,7 @@ def construct_cell_parameters_block(atoms: Atoms) -> Dict[str, Any]:
         params = dict(**cell_to_parameters(atoms.cell))
     else:
         params = {'vectors': [list(row) for row in atoms.cell[:]], 'units': 'angstrom'}
-    params['periodic'] = all(atoms.pbc)
+    params['periodic'] = list([bool(x) for x in atoms.pbc])
     return params
 
 

src/koopmans/workflows/_dft.py

@@ -36,15 +36,14 @@ def _run(self):
         if self.parameters.from_scratch:
             utils.system_call(f'rm -r {calc.parameters.outdir} 2>/dev/null', False)
 
-        calc.prefix = 'dft'
-        calc.directory = '.'
-        calc.parameters.ndr = 50
-        calc.parameters.ndw = 51
-        calc.parameters.restart_mode = 'from_scratch'
-        calc.parameters.do_orbdep = False
-        calc.parameters.fixed_state = False
-        calc.parameters.do_outerloop = True
-        calc.parameters.which_compensation = 'tcc'
+        # Setting defaults
+        defaults = {'prefix': 'dft', 'directory': '.', 'ndr': 50,
+                    'ndw': 51, 'restart_mode': 'from_scratch', 'do_orbdep': False,
+                    'fixed_state': False, 'do_outerloop': True}
+
+        for key, val in defaults.items():
+            if calc.parameters[key] is None:
+                calc.parameters[key] = val
 
         if calc.parameters.maxiter is None:
             calc.parameters.maxiter = 300
@@ -53,6 +52,7 @@ def _run(self):
             if calc.parameters.empty_states_maxstep is None:
                 calc.parameters.empty_states_maxstep = 300
 
+        # Running the calculator
         self.run_calculator(calc, enforce_ss=self.parameters.fix_spin_contamination)
 
         return calc

src/koopmans/workflows/_koopmans_dscf.py

@@ -23,7 +23,8 @@
 
 class KoopmansDSCFWorkflow(Workflow):
 
-    def __init__(self, *args, redo_smooth_dft: Optional[bool] = None, restart_from_old_ki: bool = False, **kwargs) -> None:
+    def __init__(self, *args, redo_smooth_dft: Optional[bool] = None,
+                 restart_from_old_ki: bool = False, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
         # The following two additional keywords allow for some tweaking of the workflow when running a singlepoint
@@ -38,7 +39,7 @@ def __init__(self, *args, redo_smooth_dft: Optional[bool] = None, restart_from_o
 
         # If periodic, convert the kcp calculation into a Γ-only supercell calculation
         kcp_params = self.calculator_parameters['kcp']
-        if all(self.atoms.pbc):
+        if any(self.atoms.pbc):
             spins: List[Optional[str]]
             if self.parameters.spin_polarized:
                 spins = ['up', 'down']
@@ -267,7 +268,7 @@ def _run(self) -> None:
         self.perform_final_calculations()
 
         # Postprocessing
-        if all(self.atoms.pbc):
+        if any(self.atoms.pbc):
             if self.parameters.calculate_bands in [None, True] and self.projections and self.kpoints.path is not None:
                 # Calculate interpolated band structure and DOS with UI
                 from koopmans import workflows
@@ -393,23 +394,25 @@ def perform_initialization(self) -> None:
             self.run_calculator(calc, enforce_ss=False)
 
             # Check the consistency between the PW and CP band gaps
-            pw_calc = [c for c in self.calculations if isinstance(
-                c, calculators.PWCalculator) and c.parameters.calculation == 'nscf'][-1]
-            pw_gap = pw_calc.results['lumo_energy'] - pw_calc.results['homo_energy']
-            cp_gap = calc.results['lumo_energy'] - calc.results['homo_energy']
-            if abs(pw_gap - cp_gap) > 2e-2 * pw_gap:
-                raise ValueError(f'PW and CP band gaps are not consistent: {pw_gap} {cp_gap}')
+            if calc.results['lumo_energy']:
+                pw_calc = [c for c in self.calculations if isinstance(
+                    c, calculators.PWCalculator) and c.parameters.calculation == 'nscf'][-1]
+                pw_gap = pw_calc.results['lumo_energy'] - pw_calc.results['homo_energy']
+                cp_gap = calc.results['lumo_energy'] - calc.results['homo_energy']
+                if abs(pw_gap - cp_gap) > 2e-2 * pw_gap:
+                    raise ValueError(f'PW and CP band gaps are not consistent: {pw_gap} {cp_gap}')
 
             # The CP restarting from Wannier functions must be already converged
             Eini = calc.results['convergence']['filled'][0]['Etot']
             Efin = calc.results['energy']
             if abs(Efin - Eini) > 1e-6 * abs(Efin):
-                raise ValueError(f'Too much difference between the initial and final CP energies: {Eini} {Efin}')
+                utils.warn(f'There is a large difference ({Efin - Eini:.5f} Ha) between the initial and final CP energies ({Eini:.3f} and {Efin:.3f} Ha). This should not happen.')
 
             # Add to the outdir of dft_init a link to the files containing the Wannier functions
             dst = Path(f'{calc.parameters.outdir}/{calc.parameters.prefix}_{calc.parameters.ndw}.save/K00001/')
             for file in ['evc_occupied1.dat', 'evc_occupied2.dat', 'evc0_empty1.dat', 'evc0_empty2.dat']:
-                utils.symlink(f'{restart_dir}/{file}', dst, force=True)
+                if (restart_dir / file).exists():
+                    utils.symlink(f'{restart_dir}/{file}', dst, force=True)
 
         elif self.parameters.functional in ['ki', 'pkipz']:
             calc = self.new_kcp_calculator('dft_init')
@@ -504,7 +507,8 @@ def perform_alpha_calculations(self) -> None:
             # Do a KI/KIPZ calculation with the updated alpha values
             restart_from_wannier_pwscf = True if self.parameters.init_orbitals in [
                 'mlwfs', 'projwfs'] and not self._restart_from_old_ki and i_sc == 1 else None
-            if self.parameters.task in ['trajectory', 'convergence_ml'] and self.ml.input_data_for_ml_model == 'orbital_density':
+            if self.parameters.task in ['trajectory', 'convergence_ml'] \
+                    and self.ml.input_data_for_ml_model == 'orbital_density':
                 print_real_space_density = True
             else:
                 print_real_space_density = False
@@ -642,13 +646,13 @@ def perform_alpha_calculations(self) -> None:
 
                         alpha, error = self.calculate_alpha_from_list_of_calcs(
                             calcs, trial_calc, band, filled=band.filled)
-                        
+
                         # Mixing
                         alpha = self.parameters.alpha_mixing * alpha + (1 - self.parameters.alpha_mixing) * band.alpha
 
                     warning_message = 'The computed screening parameter is {0}. Proceed with caution.'
                     failure_message = 'The computed screening parameter is significantly {0}. This should not ' \
-                         'happen. Decrease alpha_mixing and/or change alpha_guess.'
+                        'happen. Decrease alpha_mixing and/or change alpha_guess.'
 
                     if alpha < -0.1:
                         raise ValueError(failure_message.format('less than 0'))
@@ -666,7 +670,7 @@ def perform_alpha_calculations(self) -> None:
 
                 # add alpha to training data
                 if self.ml.use_ml and not use_prediction:
-                    mlfit.print_error_of_single_orbital(alpha_predicted, alpha, indent=self.print_indent+2)
+                    mlfit.print_error_of_single_orbital(alpha_predicted, alpha, indent=self.print_indent + 2)
                     mlfit.add_training_data(band)
                     # if the user wants to train on the fly, train the model after the calculation of each orbital
                     if self.ml.train_on_the_fly:
@@ -697,7 +701,8 @@ def perform_alpha_calculations(self) -> None:
         else:
             self.print('Screening parameters have been determined but are not necessarily converged')
 
-    def perform_fixed_band_calculations(self, band, trial_calc, i_sc, alpha_dep_calcs, index_empty_to_save, outdir_band, directory, alpha_indep_calcs) -> None:
+    def perform_fixed_band_calculations(self, band, trial_calc, i_sc, alpha_dep_calcs, index_empty_to_save, outdir_band,
+                                        directory, alpha_indep_calcs) -> None:
         # Perform the fixed-band-dependent calculations
         if self.parameters.functional in ['ki', 'pkipz']:
             if band.filled:
@@ -1034,11 +1039,6 @@ def new_kcp_calculator(self, calc_presets: str = 'dft_init',
         if 'print' in calc.prefix:
             calc.parameters.print_wfc_anion = True
 
-        if self.parameters.mt_correction:
-            calc.parameters.which_compensation = 'tcc'
-        else:
-            calc.parameters.which_compensation = 'none'
-
         # If we are using frozen orbitals, we override the above logic and freeze the variational orbitals
         # post-initialization
         if self.parameters.frozen_orbitals and 'init' not in calc.prefix and not any([s == calc.prefix for s in

src/koopmans/workflows/_unfold_and_interp.py

@@ -67,14 +67,25 @@ def _run(self) -> None:
         else:
             spins = [None]
 
+        energies_to_merge = []
         for spin, band_filling in zip(spins, self.bands.filling):
             # Extract the centers and spreads corresponding to this particular spin
             centers = np.array([center for c, p in zip(w90_calcs, self.projections)
                                for center in c.results['centers'] if p.spin == spin])
             spreads = np.array([spread for c, p in zip(w90_calcs, self.projections)
                                for spread in c.results['spreads'] if p.spin == spin])
 
-            for filled, filling in zip([True, False], ['occ', 'emp']):
+            energies_to_merge.append([])
+
+            filleds = [True]
+            if not all(band_filling):
+                filleds.append(False)
+
+            for filled in filleds:
+                if filled:
+                    filling = 'occ'
+                else:
+                    filling = 'emp'
                 calc_presets = filling
                 if spin:
                     calc_presets += '_' + spin
@@ -97,17 +108,18 @@ def _run(self) -> None:
                 # Run the calculator
                 self.run_calculator(calc, enforce_ss=False)
 
+                # Store the bandstructure energies that we will merge below
+                energies_to_merge[-1].append(self.calculations[-1].results['band structure'].energies)
+
         # Merge the two calculations to print out the DOS and bands
         calc = self.new_ui_calculator('merge')
 
         # Merge the bands
         if self.parameters.spin_polarized:
-            energies = [[c.results['band structure'].energies for c in subset]
-                        for subset in [self.calculations[-4:-2], self.calculations[-2:]]]
-            reference = np.max([e[0] for e in energies])
+            reference = np.max([e[0] for e in energies_to_merge])
             energies_np = np.concatenate([np.concatenate(e, axis=2) for e in energies], axis=0)
         else:
-            energies = [c.results['band structure'].energies for c in self.calculations[-2:]]
+            energies = energies_to_merge[0]
             reference = np.max(energies[0])
             energies_np = np.concatenate(energies, axis=2)
         calc.results['band structure'] = BandStructure(self.kpoints.path, energies_np, reference=reference)

src/koopmans/workflows/_workflow.py

@@ -117,7 +117,6 @@ def __init__(self, atoms: Atoms,
                  autogenerate_settings: bool = True,
                  version: Optional[str] = None,
                  **kwargs: Dict[str, Any]):
-
         # Parsing parameters
         self.parameters = settings.WorkflowSettingsDict(**parameters)
         for key, value in kwargs.items():
@@ -229,12 +228,13 @@ def __init__(self, atoms: Atoms,
 
         # Before saving the calculator_parameters, automatically generate some keywords and perform some sanity checks
         if self.parameters.task != 'ui' and autogenerate_settings:
+
             # Automatically calculate nelec/nelup/neldw/etc using information contained in the pseudopotential files
             # and the kcp settings
             nelec = nelec_from_pseudos(self.atoms, self.pseudopotentials, self.parameters.pseudo_directory)
-            tot_charge = calculator_parameters['kcp'].get('tot_charge', 0)
+            tot_charge = kwargs.get('tot_charge', calculator_parameters['kcp'].get('tot_charge', 0))
             nelec -= tot_charge
-            tot_mag = calculator_parameters['kcp'].get('tot_magnetization', nelec % 2)
+            tot_mag = kwargs.get('tot_magnetization', calculator_parameters['kcp'].get('tot_magnetization', nelec % 2))
             nelup = int(nelec / 2 + tot_mag / 2)
             neldw = int(nelec / 2 - tot_mag / 2)
 
@@ -435,11 +435,11 @@ def _run_sanity_checks(self):
             self.parameters.calculate_alpha = False
 
         # Checking periodic image correction schemes
-        if not self.parameters.calculate_alpha:
-            # If we are not calculating alpha, we do not consider charged systems and therefore we don't need image
-            # corrections, so we skip the following checks
-            pass
-        elif all(self.atoms.pbc):
+        # if not self.parameters.calculate_alpha:
+        #     # If we are not calculating alpha, we do not consider charged systems and therefore we don't need image
+        #     # corrections, so we skip the following checks
+        #     pass
+        if all(self.atoms.pbc):
             if self.parameters.method == 'dfpt':
                 # For DPFT, we use gb_correction
                 if self.parameters.gb_correction is None:
@@ -677,6 +677,21 @@ def new_calculator(self,
         # Create the calculator
         calc = calc_class(atoms=copy.deepcopy(self.atoms), **all_kwargs)
 
+        # Add Martyna-Tuckerman settings for kcp calculators
+        if calc.parameters.is_valid('which_compensation'):
+            if self.parameters.mt_correction:
+                if not any(self.atoms.pbc):
+                    calc.parameters.which_compensation = 'tcc'
+                elif all(self.atoms.pbc == [False, False, True]):
+                    calc.parameters.which_compensation = 'tcc1d'
+                elif all(self.atoms.pbc == [True, True, False]):
+                    calc.parameters.which_compensation = 'tcc2d'
+                else:
+                    raise ValueError('Martyna-Tuckerman correction not implemented for this geometry; for 1D and 2D '
+                                     'systems please use z as the unique axis')
+            else:
+                calc.parameters.which_compensation = 'none'
+
         # Add the directory if provided
         if directory is not None:
             calc.directory = directory