Improve the effective mass detection algorithm

Make the algorithm more robust - now filter by intensity first.
Also added ability to manually pass the extrema points when
computing effective masses.

easyunfold/cli.py

@@ -267,19 +267,17 @@ def wrapper(*args, **kwargs):
 @click.option('--spin', type=int, default=0, help='Index of the spin channel.', show_default=True)
 @click.option('--npoints', type=int, default=3, help='Number of kpoints used for fitting from the extrema.', show_default=True)
 @click.option('--extrema-detect-tol', type=float, default=0.01, help='Tolerance for band extrema detection.', show_default=True)
-@click.option('--degeneracy-detect-tol',
-              type=float,
-              default=0.01,
-              help='Tolerance for band degeneracy detection at extrema.',
-              show_default=True)
 @click.option('--nocc', type=int, help='DEV: Use this band as the extrema at all kpoints.')
 @click.option('--plot', is_flag=True, default=False)
 @click.option('--plot-fit', is_flag=True, default=False, help='Generate plots of the band edge and parabolic fits.')
 @click.option('--fit-label', help='Which branch to use for plot fitting. e.g. electrons:0', default='electrons:0', show_default=True)
 @click.option('--band-filter', default=None, type=int, help='Only displace information for this band.')
 @click.option('--out-file', '-o', default='unfold-effective-mass.png', help='Name of the output file.', show_default=True)
-def unfold_effective_mass(ctx, intensity_threshold, spin, band_filter, npoints, extrema_detect_tol, degeneracy_detect_tol, nocc, plot,
-                          plot_fit, fit_label, out_file):
+@click.option('--emin', type=float, default=-5., help='Minimum energy in eV relative to the reference.', show_default=True)
+@click.option('--emax', type=float, default=5., help='Maximum energy in eV relative to the reference.', show_default=True)
+@click.option('--manual-extrema', help='Manually specify the extrema to use for fitting, in the form "mode,k_index,band_index"')
+def unfold_effective_mass(ctx, intensity_threshold, spin, band_filter, npoints, extrema_detect_tol, nocc, plot, plot_fit, fit_label,
+                          out_file, emin, emax, manual_extrema):
     """
     Compute and print effective masses by tracing the unfolded weights.
 
@@ -293,7 +291,7 @@ def unfold_effective_mass(ctx, intensity_threshold, spin, band_filter, npoints,
     from easyunfold.unfold import UnfoldKSet
     from tabulate import tabulate
     unfoldset: UnfoldKSet = ctx.obj['obj']
-    efm = EffectiveMass(unfoldset, intensity_tol=intensity_threshold, extrema_tol=extrema_detect_tol, degeneracy_tol=degeneracy_detect_tol)
+    efm = EffectiveMass(unfoldset, intensity_tol=intensity_threshold, extrema_tol=extrema_detect_tol)
 
     click.echo('Band extrema data:')
     table = []
@@ -306,12 +304,20 @@ def unfold_effective_mass(ctx, intensity_threshold, spin, band_filter, npoints,
 
     if nocc:
         efm.set_nocc(nocc)
-    output = efm.get_effective_masses(ispin=spin, npoints=npoints)
+    if manual_extrema is None:
+        output = efm.get_effective_masses(ispin=spin, npoints=npoints)
+    else:
+        mode, ik, ib = manual_extrema.split(',')
+        ik = int(ik)
+        ib = int(ib)
+        click.echo(f'Using manually passed kpoint and band: {ik},{ib}')
+        output = efm.get_effective_masses(ispin=spin, npoints=npoints, mode=mode, iks=[ik], iband=[[ib]])
 
     # Filter by band if requested
     if band_filter is not None:
         for carrier in ['electrons', 'holes']:
-            output[carrier] = [entry for entry in output[carrier] if entry['band_index'] == band_filter]
+            if carrier in output:
+                output[carrier] = [entry for entry in output[carrier] if entry['band_index'] == band_filter]
 
     ## Print data
     def print_data(entries, tag='me'):
@@ -332,10 +338,10 @@ def print_data(entries, tag='me'):
         click.echo(tabulate(table, headers=['index', 'Kind', 'Effective mass', 'Band index', 'from', 'to']))
 
     click.echo('Electron effective masses:')
-    print_data(output['electrons'], 'm_e')
+    print_data(output.get('electrons', []), 'm_e')
     print('')
     click.echo('Hole effective masses:')
-    print_data(output['holes'], 'm_h')
+    print_data(output.get('holes', []), 'm_h')
 
     click.echo('Unfolded band structure can be ambiguous, please cross-check with the spectral function plot.')
 
@@ -344,7 +350,7 @@ def print_data(entries, tag='me'):
         plotter = UnfoldPlotter(unfoldset)
         click.echo('Generating spectral function plot for visualising detected branches...')
         engs, sf = unfoldset.get_spectral_function()
-        plotter.plot_effective_mass(efm, engs, sf, effective_mass_data=output, save=out_file)
+        plotter.plot_effective_mass(efm, engs, sf, effective_mass_data=output, save=out_file, ylim=(emin, emax))
 
     elif plot_fit:
         from easyunfold.plotting import UnfoldPlotter

easyunfold/effective_mass.py

@@ -8,7 +8,7 @@
 from scipy.optimize import curve_fit
 
 from .unfold import UnfoldKSet
-# pylint: disable=invalid-name
+# pylint: disable=invalid-name,too-many-locals
 
 eV_to_hartree = physical_constants['electron volt-hartree relationship'][0]
 bohr_to_m = physical_constants['Bohr radius'][0]
@@ -54,21 +54,24 @@ def f(x, alpha, d):
     # coefficient is currently in eV/Angstrom^2/h_bar^2
     # want it in atomic units so Hartree/bohr^2/h_bar^2
     eff_mass = (angstrom_to_bohr**2 / eV_to_hartree) / c
-    return eff_mass
+    return eff_mass, fit
 
 
 def fitted_band(x: np.ndarray, eff_mass: float) -> np.ndarray:
     """Return fitted effective mass curve"""
     c = (angstrom_to_bohr**2 / eV_to_hartree) / eff_mass
-    x0 = x - x[0]
-    return x0 + x[0], c / 2 * x0**2
+    return c / 2 * x**2
 
 
-def points_with_tol(array, value, tol=1e-4):
+def points_with_tol(array, value, tol=1e-4, sign=1):
     """
     Return the indices and values of points in an array close to the value with a tolerance
     """
-    idx = np.where(np.abs(array - value) < tol)[0]
+    if sign == 0:
+        diff = abs(array - value)
+    else:
+        diff = (array - value) * sign
+    idx = np.where((-1e-3 < diff) & (diff < tol))[0]
     return idx, array[idx]
 
 
@@ -80,7 +83,6 @@ def __init__(
         unfold: UnfoldKSet,
         intensity_tol: float = 1e-1,
         extrema_tol: float = 1e-3,
-        degeneracy_tol: float = 1e-2,
         parabolic: bool = True,
         npoints: float = 3,
     ):
@@ -95,7 +97,6 @@ def __init__(
         self.unfold: UnfoldKSet = unfold
         self.intensity_tol = intensity_tol
         self.extrema_detect_tol = extrema_tol
-        self.degeneracy_tol = degeneracy_tol
         self.parabolic = parabolic
         self.nocc = None  # Number of occupied bands
         if npoints < 3:
@@ -114,28 +115,27 @@ def kpoints(self):
     def kpoints_labels(self):
         return self.unfold.kpoint_labels
 
-    def get_band_extrema(self, mode: str = 'cbm', extrema_tol: float = None, degeneracy_tol: float = None, ispin=0):
+    def get_band_extrema(self, mode: str = 'cbm', extrema_tol: float = None, ispin=0):
         """
-        Obtain the kpoint idx of band maximum, sub indices in th set and the band indices.
+        Obtain the kpoint idx of band extrema, sub indices in th set and the band indices.
 
         The search takes two steps, first the kpoints at the band extrema is located by comparing the
-        band energies with that recorded in supplied *cbm* and *vbm*, based on the `exgtrema_tol`.
-        Afterwards, the band indices are selected at the these kpoints using `degeneracy_tol`.
+        band energies with that recorded in supplied *cbm* and *vbm*, based on the `extrema_tol`.
+        Afterwards, the band indices are selected at the these kpoints using
 
         Returns:
             A tuple of extrema locations including a list of kpoint indices, sub-indices within
             the set and the band indices at each kpoint that is within the `tol` set.
         """
         if extrema_tol is None:
             extrema_tol = self.extrema_detect_tol
-        if degeneracy_tol is None:
-            degeneracy_tol = self.degeneracy_tol
 
         intensity_tol = self.intensity_tol
 
         if mode not in ['cbm', 'vbm']:
             raise ValueError(f'Unknown mode {mode}')
-        cbm = self.unfold.calculated_quantities[mode]
+
+        eref = self.unfold.calculated_quantities[mode]
         weights = self.unfold.calculated_quantities['spectral_weights_per_set']
 
         # Indices of the kpoint corresponding to the CBM
@@ -144,22 +144,20 @@ def get_band_extrema(self, mode: str = 'cbm', extrema_tol: float = None, degener
         cbm_indices = []
         for ik, wset in enumerate(weights):
             for isubset in range(wset.shape[1]):
-                if np.any(np.abs(wset[ispin, isubset, :, 0] - cbm) < extrema_tol):
-                    itmp, _ = points_with_tol(wset[ispin, isubset, :, 0], cbm, extrema_tol)
-                    # Check if it has sufficient intensity
-                    if np.max(wset[ispin, isubset, itmp, 1]) < intensity_tol:
-                        continue
-                    # Select this kpoints
-                    k_indicies.append(ik)
-                    # Select all band indices within the tolerance
-                    k_subset_indices.append(isubset)
-                    # Stop looking at other kpoints in the k subset if found
-                    break
-
-        # Go through each case
-        for ik, iksub in zip(k_indicies, k_subset_indices):
-            itmp, _ = points_with_tol(weights[ik][ispin, iksub, :, 0], cbm, degeneracy_tol)
-            cbm_indices.append(itmp)
+                etmp = weights[ik][ispin, isubset, :, 0]
+                wtmp = weights[ik][ispin, isubset, :, 1]
+                # Filter by intensity
+                mask = wtmp > intensity_tol
+                midx = np.where(mask)[0]
+                # Find points close to the reference energy (vbm/cbm)
+                itmp, _ = points_with_tol(etmp[mask], eref, extrema_tol, 0)
+                if len(itmp) == 0:
+                    continue
+                # Reconstruct the valid extrema indices
+                itmp = midx[itmp]
+                cbm_indices.append(itmp)
+                k_indicies.append(ik)
+                k_subset_indices.append(isubset)
 
         return k_indicies, k_subset_indices, cbm_indices
 
@@ -191,6 +189,8 @@ def _get_fitting_data(self, kidx: int, iband: int, direction=1, ispin=0, npoints
         npoints = self.get_npoints(npoints)
         for i in range(npoints):
             idx = istart + i * direction
+            if idx >= len(dists) or idx < 0:
+                break
             kdists.append(dists[idx])
             # Get the spectral weight array
             sw = weights[idx]
@@ -200,29 +200,40 @@ def _get_fitting_data(self, kidx: int, iband: int, direction=1, ispin=0, npoints
             # Compute the effective energy weighted by intensity and kpoint weighting
             eng_effective = np.sum(engs * intensities * kw) / np.sum(intensities * kw)
             engs_effective.append(eng_effective)
-        return kdists, engs_effective
+
+        # Normalise the fitting data
+        kdists_norm = np.array(kdists) - kdists[0]
+        engs_norm = np.array(engs_effective)
+        engs_norm -= engs_norm[0]
+
+        kdists_norm = np.concatenate([-kdists_norm[::-1], kdists_norm])
+        engs_norm = np.concatenate([engs_norm[::-1], engs_norm])
+
+        return kdists_norm, engs_norm, (kdists, engs_effective)
 
     def get_npoints(self, override: Union[float, None] = None):
         """Get the number of points used for fitting"""
         if override is None:
             return self.npoints
         return override
 
-    def get_effective_masses(self, npoints: Union[float, None] = None, ispin=0):
+    def get_effective_masses(self, npoints: Union[float, None] = None, ispin=0, iks=None, iband=None, mode=None):
         """
         Workout the effective masses based on the unfolded band structure
         """
         outputs = {}
-        for mode in ['cbm', 'vbm']:
-            name = 'electrons' if mode == 'cbm' else 'holes'
-            outputs[name] = self._get_effective_masses(mode, npoints=npoints, ispin=ispin)
+        mode = ['cbm', 'vbm'] if mode is None else [mode]
+        for mname in mode:
+            name = 'electrons' if mname == 'cbm' else 'holes'
+            outputs[name] = self._get_effective_masses(mname, npoints=npoints, ispin=ispin, iband=iband, iks=iks)
         return outputs
 
-    def _get_effective_masses(self, mode: str = 'cbm', ispin: int = 0, npoints: Union[None, int] = None):
+    def _get_effective_masses(self, mode: str = 'cbm', ispin: int = 0, npoints: Union[None, int] = None, iks=None, iband=None):
         """
         Work out the effective masses based on the unfolded band structure for CBM or VBM
         """
-        iks, _, iband = self.get_band_extrema(mode=mode)
+        if iks is None or iband is None:
+            iks, _, iband = self.get_band_extrema(mode=mode)
         # Override occupations
         if self.nocc:
             iband = [self.nocc for _ in iband]
@@ -242,9 +253,9 @@ def _get_effective_masses(self, mode: str = 'cbm', ispin: int = 0, npoints: Unio
                     continue
                 # Get fitting data for each (degenerate) band at the extrema
                 for band_id in idxb:
-                    kdists, engs_effective = self._get_fitting_data(idxk, band_id, direction, ispin=ispin, npoints=npoints)
+                    kdists, engs_effective, raw_fit_values = self._get_fitting_data(idxk, band_id, direction, ispin=ispin, npoints=npoints)
 
-                    me = fit_effective_mass(kdists, engs_effective, parabolic=self.parabolic)
+                    me, fit = fit_effective_mass(kdists, engs_effective, parabolic=self.parabolic)
 
                     # If the identified edge is not in the list of high symmetry point, ignore it
                     # This mitigate the problem where the CBM can be duplicated....
@@ -262,7 +273,9 @@ def _get_effective_masses(self, mode: str = 'cbm', ispin: int = 0, npoints: Unio
                         'type': 'electrons' if mode == 'cbm' else 'holes',
                         'raw_data': {
                             'kpoint_distances': kdists,
-                            'effective_energies': engs_effective
+                            'effective_energies': engs_effective,
+                            'fit_res': fit,
+                            'raw_fit_values': raw_fit_values,
                         }
                     })
         results.sort(key=lambda x: abs(abs(x['effective_mass'])))

easyunfold/plotting.py

@@ -40,7 +40,8 @@ def __init__(self, unfold: UnfoldKSet):
         """
         self.unfold = unfold
 
-    def plot_dos(self, ax, dos_plotter, dos_label, dos_options, ylim, eref, atoms=None, colours=None, orbitals_subplots=None):
+    @staticmethod
+    def plot_dos(ax, dos_plotter, dos_label, dos_options, ylim, eref, atoms=None, colours=None, orbitals_subplots=None):
         """
         Prepare and plot the density of states.
         """
@@ -410,38 +411,40 @@ def plot_effective_mass(self,
         kvbm = eff.get_band_extrema(mode='vbm')[0]
         all_k = sorted(list(set(list(kcbm) + list(kvbm))))
         kdist = self.unfold.get_kpoint_distances()
-        xwidth = 0.2
         if eref is None:
             eref = self.unfold.calculated_quantities['vbm']
 
-        fig, axes = plt.subplots(1, len(all_k))
+        fig, axes = plt.subplots(1, len(all_k), figsize=(4 * len(all_k), 3), dpi=300, squeeze=False)
         if effective_mass_data is None:
             effective_mass_data = eff.get_effective_masses()
 
         # Plot the spectral function
-        for (ik, ax) in zip(all_k, axes):
+        xwidth = abs(kdist[1] - kdist[0])
+        for (ik, ax) in zip(all_k, axes[0]):
             self.plot_spectral_function(engs, sf, ax=ax, eref=eref, **kwargs)
             xk = kdist[ik]
             xlim = (xk - xwidth / 2, xk + xwidth / 2)
             ax.set_xlim(xlim)
             ax.set_title(f'Kpoint: {ik}')
 
-        elec = effective_mass_data['electrons']
+        elec = effective_mass_data.get('electrons', [])
         # Plot the detected effective mass fitting data on top
         for entry in elec:
             ik = entry['kpoint_index']
             iax = all_k.index(ik)
-            x = entry['raw_data']['kpoint_distances']
-            y = entry['raw_data']['effective_energies']
-            axes[iax].plot(x, np.asarray(y) - eref, '-o', color='C1')
+            x = entry['raw_data']['raw_fit_values'][0]
+            y = entry['raw_data']['raw_fit_values'][1]
+            axes[0, iax].plot(x, np.asarray(y) - eref, '-o', color='C1')
+            axes[0, iax].set_xlim(min(x) - xwidth / 2, max(x) + xwidth / 2)
 
-        hole = effective_mass_data['holes']
+        hole = effective_mass_data.get('holes', [])
         for entry in hole:
             ik = entry['kpoint_index']
             iax = all_k.index(ik)
-            x = entry['raw_data']['kpoint_distances']
-            y = entry['raw_data']['effective_energies']
-            axes[iax].plot(x, np.asarray(y) - eref, '-o', color='C2')
+            x = entry['raw_data']['raw_fit_values'][0]
+            y = entry['raw_data']['raw_fit_values'][1]
+            axes[0, iax].plot(x, np.asarray(y) - eref, '-o', color='C2')
+            axes[0, iax].set_xlim(min(x) - xwidth / 2, max(x) + xwidth / 2)
         if save:
             fig.savefig(save)
 
@@ -800,11 +803,11 @@ def plot_effective_mass_fit(efm: EffectiveMass,
         x = data['raw_data']['kpoint_distances']
         y = data['raw_data']['effective_energies']
         me = data['effective_mass']
-        x1, y1 = fitted_band(x, me)
+        y1 = fitted_band(x, me)
         if ax is None:
             fig, ax = plt.subplots(1, 1)
         ax.plot(x, y, 'x-', label='Energy ')
-        ax.plot(x1, y1 + y[0], label='fitted')
+        ax.plot(x, y1, label='fitted')
         ax.legend()
         if save:
             fig.savefig(save, dpi=dpi)