Skip to content

Commit

Permalink
Added more configurations to allow users more control of the plot. Th…
Browse files Browse the repository at this point in the history
…ey are the *_params which are input into various matplotlib functions
  • Loading branch information
lllangWV committed Aug 14, 2024
1 parent 9276237 commit fce9c10
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 54 deletions.
65 changes: 65 additions & 0 deletions pyprocar/cfg/band_structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from typing import Dict, Any, List, Optional,Tuple
from enum import Enum, auto


from pyprocar.cfg.base import PlotType, BaseConfig

class BandStructureMode(Enum):
Expand Down Expand Up @@ -117,6 +118,33 @@ class BandStructureConfig(BaseConfig):
The size of the figure (width, height) in inches.
dpi : str, optional
The resolution in dots per inch. If 'figure', use the figure's dpi value.
colorbar_tick_params : Dict[str, any], optional
The colorbar tick parameters, by default None
colorbar_label_params : Dict[str, any], optional
The colorbar label parameters, by default None
x_label_params : Dict[str, any], optional
The x label parameters, by default None
y_label_params : Dict[str, any], optional
The y label parameters, by default None
title_params : Dict[str, any], optional
The title parameters, by default None
major_x_tick_params : Dict[str, any], optional
The major x tick parameters, by default None
major_y_tick_params : Dict[str, any], optional
The major y tick parameters, by default None
minor_y_tick_params : Dict[str, any], optional
The minor y tick parameters, by default None
major_y_locator : matplotlib.ticker.Locator, optional
The major y locator, by default None
minor_y_locator : matplotlib.ticker.Locator, optional
The minor y locator, by default None
multiple_locator_y_major_value : float, optional
The major value for the multiple locator, by default None
multiple_locator_y_minor_value : float, optional
The minor value for the multiple locator, by default None
Methods
-------
Expand Down Expand Up @@ -171,6 +199,43 @@ class BandStructureConfig(BaseConfig):
figure_size: Tuple[int] = field(default_factory=lambda: (9, 6))
dpi: str = 'figure'


colorbar_tick_params: Dict[str, any] = field(default_factory=lambda: {})
colorbar_label_params: Dict[str, any] = field(default_factory=lambda: {})
x_label_params: Dict[str, any] = field(default_factory=lambda: {})
y_label_params: Dict[str, any] = field(default_factory=lambda: {})
title_params: Dict[str, any] = field(default_factory=lambda: {})
# Tick Parameters
major_y_tick_params: Dict[str, any] = field(default_factory=lambda: {
"which": "major",
"axis": "y",
"direction": "inout",
"width": 1,
"length": 5,
"labelright": False,
"right": True,
"left": True
})
minor_y_tick_params: Dict[str, any] = field(default_factory=lambda: {
"which": "minor",
"axis": "y",
"direction": "in",
"left": True,
"right": True
})
major_x_tick_params: Dict[str, any] = field(default_factory=lambda: {
"which": "major",
"axis": "x",
"direction": "in"
})
major_y_locator = None
minor_y_locator = None

multiple_locator_y_major_value:float = None
multiple_locator_y_minor_value:float = None



def __post_init__(self):
"""This method is immediately called after the object is initialized.
It is useful to validate the data and set default values.
Expand Down
115 changes: 61 additions & 54 deletions pyprocar/plotter/ebs_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -528,10 +528,7 @@ def set_xticks(self,

self.ax.set_xticks(self.x[tick_positions])
self.ax.set_xticklabels(tick_names)
self.ax.tick_params(
which='major',
axis='x',
direction='in')
self.ax.tick_params(**self.config.major_x_tick_params)

def set_yticks(self,
major:float=None,
Expand All @@ -548,48 +545,50 @@ def set_yticks(self,
interval : List[float], optional
The interval of the ticks, by default None
"""
if (major is None or minor is None):
if interval is None:
interval = (self.ebs.bands.min()-abs(self.ebs.bands.min())
* 0.1, self.ebs.bands.max()*1.1)

interval = abs(interval[1] - interval[0])
if interval < 30 and interval >= 20:
major = 5
minor = 1
elif interval < 20 and interval >= 10:
major = 4
minor = 0.5
elif interval < 10 and interval >= 5:
major = 2
minor = 0.2
elif interval < 5 and interval >= 3:
major = 1
minor = 0.1
elif interval < 3 and interval >= 1:
major = 0.5
minor = 0.1
else:
pass
if major is not None and minor is not None:
self.ax.yaxis.set_major_locator(MultipleLocator(major))
self.ax.yaxis.set_minor_locator(MultipleLocator(minor))
self.ax.tick_params(
which='major',
axis="y",
direction="inout",
width=1,
length=5,
labelright=False,
right=True,
left=True)

self.ax.tick_params(
which='minor',
axis="y",
direction="in",
left=True,
right=True)
# if (major is None or minor is None):
if interval is None:
interval = (self.ebs.bands.min()-abs(self.ebs.bands.min())
* 0.1, self.ebs.bands.max()*1.1)

interval = abs(interval[1] - interval[0])
if interval < 30 and interval >= 20:
major = 5
minor = 1
elif interval < 20 and interval >= 10:
major = 4
minor = 0.5
elif interval < 10 and interval >= 5:
major = 2
minor = 0.2
elif interval < 5 and interval >= 3:
major = 1
minor = 0.1
elif interval < 3 and interval >= 1:
major = 0.5
minor = 0.1
else:
pass


if self.config.multiple_locator_y_major_value is not None:
major = self.config.multiple_locator_y_major_value
if self.config.multiple_locator_y_minor_value is not None:
minor = self.config.multiple_locator_y_minor_value


if self.config.major_y_locator is not None or self.config.minor_y_locator is not None:
if self.config.major_y_locator is not None:
self.ax.yaxis.set_major_locator(self.config.major_y_locator)
if self.config.minor_y_locator is not None:
self.ax.yaxis.set_minor_locator(self.config.minor_y_locator)
else:
if major is not None:
self.ax.yaxis.set_major_locator(MultipleLocator(major))
if minor is not None:
self.ax.yaxis.set_minor_locator(MultipleLocator(minor))

self.ax.tick_params(**self.config.major_y_tick_params)
self.ax.tick_params(**self.config.minor_y_tick_params)

def set_xlim(self, interval:List[float]=None):
"""A method to set the x limit
Expand Down Expand Up @@ -624,7 +623,7 @@ def set_xlabel(self, label:str="K vector"):
label : str, optional
String fo the x label name, by default "K vector"
"""
self.ax.set_xlabel(label)
self.ax.set_xlabel(label, **self.config.x_label_params)

def set_ylabel(self, label:str=r"E - E$_F$ (eV)"):
"""A method to set the y label
Expand All @@ -634,7 +633,7 @@ def set_ylabel(self, label:str=r"E - E$_F$ (eV)"):
label : str, optional
String fo the y label name, by default r"E - E$ (eV)"
"""
self.ax.set_ylabel(label)
self.ax.set_ylabel(label, **self.config.y_label_params)

def set_title(self, title:str="Band Structure"):
"""A method to set the title
Expand All @@ -645,7 +644,7 @@ def set_title(self, title:str="Band Structure"):
String for the title, by default "Band Structure"
"""
if self.config.title:
self.ax.set_title(label=self.config.title)
self.ax.set_title(label=self.config.title, **self.config.title_params)

def set_colorbar_title(self, title:str=None):
"""A method to set the title of the color bar
Expand All @@ -659,11 +658,19 @@ def set_colorbar_title(self, title:str=None):
title=title
else:
title=self.config.colorbar_title
self.cb.ax.tick_params(labelsize=self.config.colorbar_tick_labelsize)
self.cb.set_label(title,
size=self.config.colorbar_title_size,
rotation=270,
labelpad=self.config.colorbar_title_padding)

if self.config.colorbar_tick_params:
self.cb.ax.tick_params(**self.config.colorbar_tick_params)
else:
self.cb.ax.tick_params(labelsize=self.config.colorbar_tick_labelsize)

if self.config.colorbar_label_params:
self.cb.set_label(title, **self.config.colorbar_label_params)
else:
self.cb.set_label(title,
size=self.config.colorbar_title_size,
rotation=270,
labelpad=self.config.colorbar_title_padding)

def legend(self, labels:List[str]=None):
"""A methdo to plot the legend
Expand Down

0 comments on commit fce9c10

Please sign in to comment.