Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

More general plot class for arbitrary geometry, not only sp2 #121

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions hubbard/plot/bonds.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,16 @@ class BondOrder(GeometryPlot):
Mean-field Hubbard Hamiltonian
"""

def __init__(self, HH, selection=None, **kwargs):
def __init__(self, HH, selection=None, collection='sp2', **kwargs):

if 'cmap' not in kwargs:
kwargs['cmap'] = plt.cm.bwr

if selection is None:
super().__init__(HH.geometry, **kwargs)
super().__init__(HH.geometry, collection=collection, **kwargs)
else:
print("doing sub")
super().__init__(HH.geometry.sub(selection), **kwargs)
super().__init__(HH.geometry.sub(selection), collection=collection, **kwargs)

bbox_props = dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9)

Expand Down
12 changes: 6 additions & 6 deletions hubbard/plot/charge.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ class Charge(GeometryPlot):

"""

def __init__(self, HH, ext_geom=None, spin=[0, 1], realspace=False, **kwargs):
def __init__(self, HH, ext_geom=None, collection='sp2', bdx=2, spin=[0, 1], realspace=False, **kwargs):
# Set default kwargs
if realspace:
if 'facecolor' not in kwargs:
Expand All @@ -46,7 +46,7 @@ def __init__(self, HH, ext_geom=None, spin=[0, 1], realspace=False, **kwargs):
if 'label' not in kwargs:
kwargs['label']=r'$Q_\uparrow+Q_\downarrow$ ($e$)'

super().__init__(HH.geometry, ext_geom=ext_geom, **kwargs)
super().__init__(HH.geometry, ext_geom=ext_geom, collection=collection, bdx=bdx, **kwargs)

# Compute total charge on each site
if not isinstance(spin, list):
Expand Down Expand Up @@ -107,7 +107,7 @@ class ChargeDifference(GeometryPlot):
If True either a `sisl.SuperCell` (`sc` kwarg) or the `z` kwarg to slice the real space grid at the desired z coordinate needs to be passed
"""

def __init__(self, HH, ext_geom=None, realspace=False, **kwargs):
def __init__(self, HH, ext_geom=None, collection='sp2', bdx=2, realspace=False, **kwargs):

# Set default kwargs
if realspace:
Expand All @@ -123,7 +123,7 @@ def __init__(self, HH, ext_geom=None, realspace=False, **kwargs):
if 'label' not in kwargs:
kwargs['label']=r'$Q_\uparrow+Q_\downarrow-Q_\mathrm{NA}$ ($e$)'

super().__init__(HH.geometry, ext_geom=ext_geom, **kwargs)
super().__init__(HH.geometry, ext_geom=ext_geom, collection=collection, bdx=bdx, **kwargs)

# Compute total charge on each site, subtract neutral atom charge
chg = np.zeros((HH.geometry.na))
Expand Down Expand Up @@ -186,7 +186,7 @@ class SpinPolarization(GeometryPlot):
If True either a `sisl.SuperCell` (`sc` kwarg) or the `z` kwarg to slice the real space grid at the desired z coordinate needs to be passed
"""

def __init__(self, HH, ext_geom=None, realspace=False, **kwargs):
def __init__(self, HH, ext_geom=None, collection='sp2', bdx=2, realspace=False, **kwargs):

# Set default kwargs
if realspace:
Expand All @@ -202,7 +202,7 @@ def __init__(self, HH, ext_geom=None, realspace=False, **kwargs):
if 'label' not in kwargs:
kwargs['label']=r'$Q_\uparrow-Q_\downarrow$ ($e$)'

super().__init__(HH.geometry, ext_geom=ext_geom, **kwargs)
super().__init__(HH.geometry, ext_geom=ext_geom, collection=collection, bdx=bdx, **kwargs)

# Sum over all orbitals
chg = np.zeros((HH.geometry.na))
Expand Down
57 changes: 36 additions & 21 deletions hubbard/plot/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,9 +149,13 @@ class GeometryPlot(Plot):
in a pi-tight-binding or mean-feild Hubbard calculation e.g., as Hydrogen atoms
bdx: int, optional
added space between geometry and figure axes
collection: str or list, optional
if ``collection='sp2'`` it uses the inner routines to plot the geometry as an sp2 carbon system (default)
if ``collection=None`` it plots all atoms equally with the same cyrcle radius
alternatively `collection` can be an externally defined list of matplotlib patches
"""

def __init__(self, geometry, ext_geom=None, bdx=2, **kwargs):
def __init__(self, geometry, ext_geom=None, bdx=2, collection='sp2', **kwargs):

super().__init__(**kwargs)

Expand All @@ -173,28 +177,39 @@ def __init__(self, geometry, ext_geom=None, bdx=2, **kwargs):
g = ext_geom
else:
g = self.geometry
for ia in g:
idx = g.close(ia, R=[0.1, 1.6])
if g.atoms[ia].Z == 1: # H
aux.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.4))
elif g.atoms[ia].Z == 5: # B
self.axes.add_patch(patches.Circle((np.average(g.xyz[ia, 0]), np.average(g.xyz[ia, 1])), radius=0.7, color='r', lw=2, fill=False))
pi.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
elif g.atoms[ia].Z == 6: # C
if len(idx[1]) == 4:
# If the C atom has 4 neighbours (sp3 configuration) it will be represented
# as an aux site
aux.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
# Add a blue patch at the H positions
Hsp3 = [i for i in idx[1] if g.atoms[i].Z == 1]
self.axes.add_patch(patches.Circle((np.average(g.xyz[Hsp3, 0]), np.average(g.xyz[Hsp3, 1])), radius=1.4, alpha=0.15, fc='c'))
else:

if collection == 'sp2':
for ia in g:
idx = g.close(ia, R=[0.1, 1.6])
if g.atoms[ia].Z == 1: # H
aux.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.4))
elif g.atoms[ia].Z == 5: # B
self.axes.add_patch(patches.Circle((np.average(g.xyz[ia, 0]), np.average(g.xyz[ia, 1])), radius=0.7, color='r', lw=2, fill=False))
pi.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
elif g.atoms[ia].Z == 6: # C
if len(idx[1]) == 4:
# If the C atom has 4 neighbours (sp3 configuration) it will be represented
# as an aux site
aux.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
# Add a blue patch at the H positions
Hsp3 = [i for i in idx[1] if g.atoms[i].Z == 1]
self.axes.add_patch(patches.Circle((np.average(g.xyz[Hsp3, 0]), np.average(g.xyz[Hsp3, 1])), radius=1.4, alpha=0.15, fc='c'))
else:
pi.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
elif g.atoms[ia].Z == 7: # N
self.axes.add_patch(patches.Circle((np.average(g.xyz[ia, 0]), np.average(g.xyz[ia, 1])), radius=0.7, color='g', lw=2, fill=False))
pi.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
elif g.atoms[ia].Z == 7: # N
self.axes.add_patch(patches.Circle((np.average(g.xyz[ia, 0]), np.average(g.xyz[ia, 1])), radius=0.7, color='g', lw=2, fill=False))
elif g.atoms[ia].Z > 10: # Some other atom
aux.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.2))

elif collection is None:
# Create same patches for all atoms
for ia in g:
pi.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.7))
elif g.atoms[ia].Z > 10: # Some other atom
aux.append(patches.Circle((g.xyz[ia, 0], g.xyz[ia, 1]), radius=0.2))

elif all(isinstance(p, int) for p in collection):
# In this case collection should be a list of matplotlib.patches
pi = collection

# Pi sites
ppi = PatchCollection(pi, alpha=1., lw=1.2, edgecolor='0.6', **kw)
Expand Down
8 changes: 4 additions & 4 deletions hubbard/plot/spectrum.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ class LDOS_from_eigenstate(GeometryPlot):
In this case either a `sisl.SuperCell` (`sc` kwarg) or the `z` kwarg to slice the real space grid at the desired z coordinate needs to be passed
"""

def __init__(self, HH, wavefunction, sites=[], ext_geom=None, realspace=False, **kwargs):
def __init__(self, HH, wavefunction, sites=[], ext_geom=None, realspace=False, collection='sp2', **kwargs):

# Set default kwargs
if realspace:
Expand All @@ -182,7 +182,7 @@ def __init__(self, HH, wavefunction, sites=[], ext_geom=None, realspace=False, *
if 'cmap' not in kwargs:
kwargs['cmap'] = plt.cm.bwr

super().__init__(HH.geometry, ext_geom=ext_geom, **kwargs)
super().__init__(HH.geometry, ext_geom=ext_geom, collection=collection, **kwargs)

x = HH.geometry[:, 0]
y = HH.geometry[:, 1]
Expand Down Expand Up @@ -262,7 +262,7 @@ class LDOS(GeometryPlot):
sisl.physics.electron.PDOS: sisl method to obtain PDOS
"""

def __init__(self, HH, E, sites=[], spin=[0,1], ext_geom=None, realspace=False, eta=1e-3, distribution='lorentzian', **kwargs):
def __init__(self, HH, E, sites=[], spin=[0,1], ext_geom=None, realspace=False, eta=1e-3, distribution='lorentzian', collection='sp2', **kwargs):

# Set default kwargs
if realspace:
Expand All @@ -274,7 +274,7 @@ def __init__(self, HH, E, sites=[], spin=[0,1], ext_geom=None, realspace=False,
if 'cmap' not in kwargs:
kwargs['cmap'] = plt.cm.bwr

super().__init__(HH.geometry, ext_geom=ext_geom, **kwargs)
super().__init__(HH.geometry, ext_geom=ext_geom, collection=collection, **kwargs)

x = HH.geometry[:, 0]
y = HH.geometry[:, 1]
Expand Down
4 changes: 2 additions & 2 deletions hubbard/plot/wavefunction.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ class Wavefunction(GeometryPlot):
of the coefficient of `wf` on the atomic sites
"""

def __init__(self, HH, wf, ext_geom=None, cb_label=r'Phase', realspace=False, **kwargs):
def __init__(self, HH, wf, ext_geom=None, cb_label=r'Phase', realspace=False, collection='sp2', **kwargs):

# Set default kwargs
if realspace:
Expand All @@ -40,7 +40,7 @@ def __init__(self, HH, wf, ext_geom=None, cb_label=r'Phase', realspace=False, **
if 'cmap' not in kwargs:
kwargs['cmap'] = plt.cm.bwr

super().__init__(HH.geometry, ext_geom=ext_geom, **kwargs)
super().__init__(HH.geometry, ext_geom=ext_geom, collection=collection, **kwargs)

self.plot_wf(HH, wf, cb_label=cb_label, realspace=realspace, **kwargs)

Expand Down