elphondb.py 3d plots for the gkkp added#92
elphondb.py 3d plots for the gkkp added#92Petru-Milev wants to merge 1 commit intoyambo-code:masterfrom
Conversation
There was a problem hiding this comment.
Pull request overview
This pull request adds 3D plotting capability to the YamboElectronPhononDB class, enabling visualization of electron-phonon coupling data on 3D k-grids. The implementation preserves backward compatibility by keeping the original 2D plotting behavior as the default and adding an optional threeD parameter to enable the new functionality.
Changes:
- Added
red_kpointsattribute to store reduced k-point coordinates - Renamed the original
plot_elphmethod toplot_elph_oldfor backward compatibility - Implemented new
plot_elphmethod that supports both 2D (default) and 3D k-grid plotting via athreeDparameter
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| if plt_show: plt.show() | ||
| else: print("Plot ready.\nYou can customise adding savefig, title, labels, text, show, etc...") | ||
|
|
There was a problem hiding this comment.
The new plot_elph function is missing the @add_fig_kwargs decorator that was present in the original plot_elph_old function. This decorator provides important functionality for handling common figure kwargs like 'title', 'show', and 'savefig'. Without it, users who pass these arguments will experience unexpected behavior as the kwargs won't be processed.
| @add_fig_kwargs |
| import numpy as np | ||
| import matplotlib.pyplot as plt | ||
| import math |
There was a problem hiding this comment.
These imports are redundant since numpy is already imported as 'np' at the top of the file (line 9), matplotlib.pyplot is already imported as 'plt' (line 12), and math.ceil is available from the math module already imported (line 8). Importing inside a function is generally discouraged unless needed for lazy loading or optional dependencies. Remove these imports and use the module-level imports instead.
| (i) in k-space of the quantity A_{k}(iq,inu,ib1,ib2). | ||
| (ii) in q-space of the quantity A_{q}(ik,inu,ib1,ib2). |
There was a problem hiding this comment.
The docstring has inconsistent indentation. Line 290 has extra leading spaces before "(i)" and line 291 has extra leading spaces before "(ii)". This should be aligned consistently with the rest of the docstring (as seen in the original plot_elph_old function).
| (i) in k-space of the quantity A_{k}(iq,inu,ib1,ib2). | |
| (ii) in q-space of the quantity A_{q}(ik,inu,ib1,ib2). | |
| (i) in k-space of the quantity A_{k}(iq,inu,ib1,ib2). | |
| (ii) in q-space of the quantity A_{q}(ik,inu,ib1,ib2). |
| if plt_cbar: | ||
| if 'cmap' in kwargs.keys(): color_map = plt.get_cmap(kwargs['cmap']) | ||
| else: color_map = plt.get_cmap('viridis') | ||
|
|
There was a problem hiding this comment.
The variable 'color_map' is assigned but never used in the 2D plotting path. In the original plot_elph_old function (lines 268-269), this variable is also assigned but not used. Consider removing this dead code or using it to set the colormap explicitly for the scatter plot.
| if plt_cbar: | |
| if 'cmap' in kwargs.keys(): color_map = plt.get_cmap(kwargs['cmap']) | |
| else: color_map = plt.get_cmap('viridis') |
| ax.set_title(f"{axis} ≈ {layer_val:g} (rounded {decimals} dp)") | ||
|
|
||
| # select points in this layer | ||
| mask = (kpts_r == layer_val) |
There was a problem hiding this comment.
The masking operation uses kpts_r (rounded values) to select points but then uses the mask to index the original kpts array. This is correct behavior for selecting the original k-points that belong to a layer. However, consider adding a check to ensure that at least one point exists in each layer before attempting to plot it, to avoid potential errors with empty arrays.
| mask = (kpts_r == layer_val) | |
| mask = (kpts_r == layer_val) | |
| # If no k-points fall into this layer, skip plotting to avoid empty-array issues | |
| if not np.any(mask): | |
| continue |
| if plt_show: | ||
| plt.show() | ||
| else: | ||
| print("3D subplot grid ready.\nYou can customise adding savefig, title, labels, text, show, etc...") |
There was a problem hiding this comment.
The 3D plotting path doesn't return self.fig, while the 2D path has an explicit return statement at line 371. For API consistency and to match the expected behavior of the @add_fig_kwargs decorator (which should be added), this function should return self.fig at the end of the 3D path, similar to how the original plot_elph_old relies on the decorator to return the figure.
| self.alat = lattice.alat | ||
| self.rlat = lattice.rlat | ||
| self.car_kpoints = lattice.car_kpoints | ||
| self.red_kpoints = lattice.red_kpoints |
There was a problem hiding this comment.
The red_kpoints attribute is assigned from lattice.red_kpoints but is never used anywhere in the code. If this was intended for the 3D plotting functionality, it should be utilized, or the line should be removed if it's not needed. Consider clarifying whether this is intended for future use or if it should be removed.
| self.red_kpoints = lattice.red_kpoints |
| proj_label = f"{axis_names[xidx]}{axis_names[yidx]}" | ||
|
|
||
| nlay = len(uniq_layers) | ||
| nrows = int(math.ceil(nlay / ncols)) |
There was a problem hiding this comment.
If no unique layers are found (e.g., all k-points have the same value along the selected axis), nlay will be 0, leading to nrows = 0. Calling plt.subplots(0, ncols) will raise a ValueError. Consider adding a check before creating subplots to ensure nlay > 0, and either raise a meaningful error or fall back to 2D plotting.
| nrows = int(math.ceil(nlay / ncols)) | |
| if nlay == 0: | |
| raise ValueError( | |
| "No unique layers found along the selected axis for 3D projection. " | |
| "Try adjusting the tolerance or use the 2D plotting mode (threeD=False)." | |
| ) | |
| nrows = int(np.ceil(float(nlay) / ncols)) |
| else: print("Plot ready.\nYou can customise adding savefig, title, labels, text, show, etc...") | ||
|
|
||
| def plot_elph(self, data, kcoords=None, plt_show=False, plt_cbar=False, | ||
| threeD=False, axis='z', tol=1e-6, ncols=3, **kwargs): |
There was a problem hiding this comment.
The parameter name 'threeD' uses camelCase, which is inconsistent with the Python naming convention (PEP 8) and other parameters in this function (plt_show, plt_cbar, kcoords) which use snake_case. Consider renaming to 'three_d' or 'is_3d' for consistency with codebase conventions.
| - if plt_cbar colorbar is shown | ||
| - kwargs example: marker='H', s=300, cmap='viridis', etc. | ||
|
|
||
| NB: So far requires a 2D system. |
There was a problem hiding this comment.
The note "NB: So far requires a 2D system." is outdated since the function now supports 3D systems when threeD=True. Consider updating this to clarify that 2D systems are required only when threeD=False, or remove this note entirely since the next sentence explains the 3D functionality.
| NB: So far requires a 2D system. |
1 participant
I have added plotting of elph databases for the case of 3D k-grids. It will generate a file with subplots for each of the "layers" of the selected axis. The old 2D plotting is preserved.
The call to the function is done by
yelph.plot_elph(g_of_x, s=s, threeD=True, axis='z', plt_cbar=False, marker='H', cmap=cmap)
I tested it on Python 3.11.11, yambopy v0.6, on Tirant cluster (Univ. of Valencia cluster). The material was MoS2, with a grid of 12x12x2.
The code was generated with the assist of chatgpt, but I checked it to not introduce any errors.
Example of the output is attached to this request.
(additionally i have similar scripts for plotting the ibz/bz for the case of k-grid in 3D, but they work more as wrappers to the original 2D plotting functions, if needed I could submit them as well).
gamma_phonon_8_band_13_14_q.pdf
gamma_phonon_4_band_13_14_k.pdf