From 326bef1fb4ebab5283544be20ec531ea9810217f Mon Sep 17 00:00:00 2001 From: Matteo Giantomassi Date: Fri, 7 Feb 2025 01:11:18 +0100 Subject: [PATCH] Saving work related to OrbMag --- abipy/electrons/orbmag.py | 137 +++++++++++++++++++++++++------------- 1 file changed, 92 insertions(+), 45 deletions(-) diff --git a/abipy/electrons/orbmag.py b/abipy/electrons/orbmag.py index 5a9b2edc4..bbf1eb33b 100644 --- a/abipy/electrons/orbmag.py +++ b/abipy/electrons/orbmag.py @@ -16,16 +16,46 @@ from abipy.electrons.ebands import ElectronBands, ElectronsReader #from abipy.tools.numtools import gaussian from abipy.tools.typing import Figure -from abipy.tools.plotting import set_axlims, get_axarray_fig_plt, add_fig_kwargs +from abipy.tools.plotting import set_axlims, get_ax_fig_plt, get_axarray_fig_plt, add_fig_kwargs +#from abipy.tools.plotting import (set_axlims, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, +# get_ax3d_fig_plt, rotate_ticklabels, set_visible, plot_unit_cell, set_ax_xylabels, get_figs_plotly, +# get_fig_plotly, add_plotly_fig_kwargs, PlotlyRowColDesc, plotly_klabels, plotly_set_lims) + +def print_options_decorator(**kwargs): + """ + A decorator to temporarily set np.printoptions inside a function. + Use the decorator to apply print settings inside the function. + + @print_options_decorator(precision=2, suppress=True) + def print_array(): + print(np.array([1.234567, 7.891234, 3.456789])) + """ + def decorator(func): + import functools + @functools.wraps(func) + def wrapper(*args, **kwargs_inner): + with np.printoptions(**kwargs): + return func(*args, **kwargs_inner) + return wrapper + return decorator + class OrbmagAnalyzer: + """ + """ def __init__(self, filepaths): self.orb_files = [OrbmagFile(path) for path in filepaths] + # Consistency check + structure0 = self.orb_files[0].structure + if any(orb_file.structure != structure0 for orb_file in self.orb_files[1:]): + raise RuntimeError("ORBMAG.nc files have different structures") + # This piece of code is taken from merge_orbmag_mesh - # the main difference is that ncroots[0] is replaced by the ElectronsReader of the first OrbmagFile + # the main difference is that here ncroots[0] is replaced by + # the reader instance of the first OrbmagFile. r0 = self.orb_files[0].r mband = r0.read_dimvalue('mband') @@ -57,13 +87,12 @@ def __init__(self, filepaths): omtmp = np.matmul(rprimd, r.read_variable('orbmag_mesh')[iterm,0:ndir,isppol,ikpt,iband]) self.orbmag_merge_mesh[iterm,idir,0:ndir,isppol,ikpt,iband] = omtmp - # This piace of code has been taken from orbmag_sigij_mesh wtk = r0.read_value('kpoint_weights') occ = r0.read_value('occupations') orbmag_nterms = r0.read_dimvalue('orbmag_nterms') - self.orbmag_merge_sigij_mesh = np.zeros((orbmag_nterms,nsppol,nkpt,mband,ndir,ndir)) + self.orbmag_merge_sigij_mesh = np.zeros((orbmag_nterms, nsppol, nkpt, mband, ndir, ndir)) for iband in range(mband): for isppol in range(nsppol): @@ -77,52 +106,84 @@ def __init__(self, filepaths): self.orbmag_merge_sigij_mesh[iterm,isppol,ikpt,iband,0:ndir,0:ndir]= \ ucvol*trnrm*self.orbmag_merge_mesh[iterm,0:ndir,0:ndir,isppol,ikpt,iband] - def report_eigvals(self, report_type): + #def __str__(self) -> str: + # """String representation""" + # return self.to_string() + #def to_string(self, verbose: int = 0) -> str: + # lines = []; app = lines.append + # return "\n".join(lines) + + #@print_options_decorator(precision=2, suppress=True) + def report_eigvals(self, report_type): + """ + """ np.set_printoptions(precision=2) - terms=['CC ','VV1 ','VV2 ','NL ','LR ','A0An '] + terms = ['CC ','VV1 ','VV2 ','NL ','LR ','A0An '] orbmag_merge_sigij_mesh = self.orbmag_merge_sigij_mesh - total_sigij=orbmag_merge_sigij_mesh.sum(axis=(0,1,2,3)) - eigenvalues=-1.0E6*np.real(eigvals(total_sigij)) - isotropic=eigenvalues.sum()/3.0 - span=eigenvalues.max()-eigenvalues.min() - skew=3.0*(eigenvalues.sum()-eigenvalues.max()-eigenvalues.min()-isotropic)/span + total_sigij = orbmag_merge_sigij_mesh.sum(axis=(0,1,2,3)) + eigenvalues = -1.0E6 * np.real(eigvals(total_sigij)) + isotropic = eigenvalues.sum() / 3.0 + span = eigenvalues.max() - eigenvalues.min() + skew = 3.0 * (eigenvalues.sum() - eigenvalues.max() -eigenvalues.min() -isotropic) / span - print('\nShielding tensor eigenvalues, ppm : ',eigenvalues) - print('Shielding tensor iso, span, skew, ppm : %6.2f %6.2f %6.2f \n'%(isotropic,span,skew)) + print('\nShielding tensor eigenvalues, ppm : ', eigenvalues) + print('Shielding tensor iso, span, skew, ppm : %6.2f %6.2f %6.2f \n' % (isotropic,span,skew)) if report_type == 'T': print('Term totals') - term_sigij=orbmag_merge_sigij_mesh.sum(axis=(1,2,3)) - for iterm in range(np.size(orbmag_merge_sigij_mesh,axis=0)): - eigenvalues=-1.0E6*np.real(eigvals(term_sigij[iterm])) - print(terms[iterm]+': ',eigenvalues) + term_sigij = orbmag_merge_sigij_mesh.sum(axis=(1, 2, 3)) + for iterm in range(np.size(orbmag_merge_sigij_mesh, axis=0)): + eigenvalues = -1.0E6 * np.real(eigvals(term_sigij[iterm])) + print(terms[iterm] + ': ', eigenvalues) print('\n') - if report_type == 'B': + elif report_type == 'B': print('Band totals') - band_sigij=orbmag_merge_sigij_mesh.sum(axis=(0,1,2)) - for iband in range(np.size(orbmag_merge_sigij_mesh,axis=3)): - eigenvalues=-1.0E6*np.real(eigvals(band_sigij[iband])) - print('band '+str(iband)+' : ',eigenvalues) + band_sigij = orbmag_merge_sigij_mesh.sum(axis=(0, 1, 2)) + for iband in range(np.size(orbmag_merge_sigij_mesh, axis=3)): + eigenvalues = -1.0E6 * np.real(eigvals(band_sigij[iband])) + print('band ' + str(iband) + ' : ', eigenvalues) print('\n') - if report_type == 'TB': + elif report_type == 'TB': print('Terms in each band') - tband_sigij=orbmag_merge_sigij_mesh.sum(axis=(1,2)) - for iband in range(np.size(orbmag_merge_sigij_mesh,axis=3)): - print('band '+str(iband)+' : ') - for iterm in range(np.size(orbmag_merge_sigij_mesh,axis=0)): - eigenvalues=-1.0E6*np.real(eigvals(tband_sigij[iterm,iband])) - print(' '+terms[iterm]+': ',eigenvalues) + tband_sigij = orbmag_merge_sigij_mesh.sum(axis=(1,2)) + for iband in range(np.size(orbmag_merge_sigij_mesh, axis=3)): + print('band ' + str(iband) + ' : ') + for iterm in range(np.size(orbmag_merge_sigij_mesh, axis=0)): + eigenvalues = -1.0E6 * np.real(eigvals(tband_sigij[iterm, iband])) + print(' ' + terms[iterm] + ': ', eigenvalues) print('\n') print('\n') + @add_fig_kwargs + def plot_fatbands(self, ebands_kpath, ax=None, **kwargs) -> Figure: + """ + Plot fatbands. + Args: + ebands_kpath + """ + ebands_kpath = ElectronBands.as_ebands(ebands_kpath) + # TODO: Use SKW to interpolate ... + # I'd start by weighting each band and kpt by trace(sigij)/3.0, the isotropic part of sigij, + # both as a term-by-term plot and as a single weighting produced by summing over all 6 terms. + ebands_kpath.plot(ax=ax, show=False) + + #from abipy.core.skw import SkwInterpolator + #skw = SkwInterpolator(lpratio, kpts, eigens, fermie, nelect, cell, symrel, has_timrev, + # filter_params=None, verbose=1) + #print(skw) + #skw.eval_sk(spin, kpt, der1=None, der2=None) -> np.ndarray: + + ax, fig, plt = get_ax_fig_plt(ax=ax) + + return fig class OrbmagFile(AbinitNcFile, Has_Header, Has_Structure, Has_ElectronBands): @@ -133,13 +194,13 @@ class OrbmagFile(AbinitNcFile, Has_Header, Has_Structure, Has_ElectronBands): """ @classmethod - def from_file(cls, filepath: str) -> OrbMagFile: + def from_file(cls, filepath: str) -> OrbmagFile: """Initialize the object from a netcdf_ file""" return cls(filepath) def __init__(self, filepath: str): super().__init__(filepath) - self.r = r = ElectronsReader(filepath) + self.r = ElectronsReader(filepath) # Initialize the electron bands from file self.natom = len(self.structure) @@ -185,20 +246,6 @@ def to_string(self, verbose: int = 0) -> str: app("nsppol: %d, nkpt: %d, mband: %d" % (self.nsppol, self.nkpt, self.mband)) app("") - #if self.prtdos == 3: - # # Print table with info on atoms. - # table = [["Idx", "Symbol", "Reduced_Coords", "Lmax", "Ratsph [Bohr]", "Has_Atom"]] - # for iatom, site in enumerate(self.structure): - # table.append([ - # iatom, - # site.specie.symbol, - # "%.5f %.5f %.5f" % tuple(site.frac_coords), - # self.lmax_atom[iatom], - # self.ratsph_type[self.typat[iatom]], - # "Yes" if self.has_atom[iatom] else "No", - # ]) - # app(tabulate(table, headers="firstrow")) - if verbose > 1: app("") app(self.hdr.to_str(verbose=verbose, title="Abinit Header"))