Merge pull request #87 from perturbo-code/ivan-spectroscopy

Ultrafast spectroscopy: pump pulse generation

.gitignore

@@ -133,3 +133,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# vscode
+.vscode

docs/perturbo_docs/.gitignore

@@ -0,0 +1,4 @@
+*.f90
+*.html
+*.png
+*.svg

docs/perturbo_docs/input_parameters.py

@@ -37,6 +37,16 @@ def main():
     # title of the html page
     title = "Quantum Espresso to PERTURBO input parameters"
 
+    print("THE input_examples.html FILE HAS TO BE COPIED TO THE PERTURBO WEBSITE DIRECTORY:")
+    print("<repo path>/perturbo-code.github.io/_includes/\n")
+
+    print("OTHER HTML FILES HAVE TO BE COPIED MANUALLY TO THE PERTURBO WEBSITE DIRECTORY:")
+    print("<repo path>/perturbo-code.github.io/pages/mydoc/\n")
+
+    print("F90 FILES HAVE TO BE COPIED MANUALLY TO THE PERTURBO SOURCE CODE DIRECTORY:")
+    print("<repo path>/perturbo/pert-src/ and")
+    print("<repo path>/perturbo/qe2pert/\n")
+
     print("Generating " + fout_name + " ...")
     htmltools.create_html_file(
         yaml_dict, input_data, html_path, fout_name, title, "qe2pert"

setup.cfg

@@ -1,3 +1,3 @@
 [flake8]
 max-line-length = 160
-ignore = E114, E127, E124, E221, W293, E721, W503, F841, F401
+ignore = E114, E127, E124, E221, W293, E721, W503, F841, F401, E202

src/perturbopy/postproc/__init__.py

@@ -17,4 +17,17 @@
 from .dbs.units_dict import UnitsDict
 from .dbs.recip_pt_db import RecipPtDB
 
-from .utils import constants, plot_tools, lattice
+from .utils import constants, plot_tools, lattice, spectra_generate_pulse, timing
+
+import warnings
+import os
+
+
+# Define a custom format for warnings
+def custom_formatwarning(message, category, filename, lineno, line=None):
+    basename = os.path.basename(filename)
+    return f"{category.__name__} in {basename} at line {lineno}: {message}\n"
+
+
+# Apply the custom format
+warnings.formatwarning = custom_formatwarning

src/perturbopy/postproc/calc_modes/dyna_run.py

@@ -45,12 +45,15 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
         kpoint = np.array(tet_file['kpts_all_crys_coord'][()])
 
         self.kpt = RecipPtDB.from_lattice(kpoint, "crystal", self.lat, self.recip_lat)
-        self.bands = UnitsDict.from_dict
 
         energies = np.array(cdyna_file['band_structure_ryd'][()])
         energies_dict = {i + 1: np.array(energies[:, i]) for i in range(0, energies.shape[1])}
         self.bands = UnitsDict.from_dict(energies_dict, 'Ry')
 
+        # Raw arrays
+        self._kpoints = kpoint
+        self._energies = energies
+
         self._data = {}
 
         self.num_runs = cdyna_file['num_runs'][()]
@@ -69,7 +72,7 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
                 snap_t = np.zeros((numb, numk, num_steps), dtype=np.float64)
 
                 for itime in range(num_steps):
-                    snap_t[:, :, itime] = cdyna_file[dyn_str][f'snap_t_{itime+1}'][()].T
+                    snap_t[:, :, itime] = cdyna_file[dyn_str][f'snap_t_{itime + 1}'][()].T
 
                 # Get E-field, which is only present if nonzero
                 if "efield" in cdyna_file[dyn_str].keys():

src/perturbopy/postproc/calc_modes/spectral_cumulant.py

@@ -121,7 +121,7 @@ def plot_Aw(self, ax, ik=0, it=0, ib=0):
         # normalize
         A0w = A0w / (np.sum(A0w) * freq_step)
         # plot
-        ax.plot(freq_array, A0w, lw=2, label=f'T={ int(self.temp_array[it])} K')
+        ax.plot(freq_array, A0w, lw=2, label=f'T={int(self.temp_array[it])} K')
         ax.legend(fontsize=18)
         plt.ylim([0, np.max(A0w) * 1.1])
         plt.xlabel(r'$\omega-\epsilon_{nk}$ (eV)', fontsize=20)

src/perturbopy/postproc/utils/constants.py

@@ -14,7 +14,7 @@
                       'J': ['j', 'joule', 'joules'],
                       'Ry': ['ry', 'rydberg', 'rydbergs']}
 
-energy_units_vals = {'eV': (2.7211396132, 1), 'Ha': (1, 0), 'J': (4.359748199, -18), 'Ry': (0.5, 0)}
+energy_units_vals = {'eV': (2.7211396132, 1), 'Ha': (1, 0), 'J': (4.359748199, -18), 'Ry': (2.0, 0)}
 
 length_units_names = {'bohr': ['bohr', 'a.u', 'atomic units', 'au'], 'angstrom': ['angstrom, a'], 'm': ['m', 'meter']}
 
@@ -215,7 +215,7 @@ def conversion_factor(init_units, final_units, units_names, units_vals):
         >>> conversion_factor('a.u', 'angstrom', {'bohr': ['bohr', 'a.u'], 'angstrom': ['angstrom, a']},
                                                  {'bohr': (1, 0), 'angstrom': (0.529177249, 0)})
         0.529177249
-         
+
     """
 
     init_prefix, init_units = find_prefix_and_base_units(init_units, units_names)
@@ -256,7 +256,7 @@ def energy_conversion_factor(init_units, final_units):
     3.674930882447527e-05
     >>> energy_conversion_factor('Ry', 'Ha')
     2.0
-       
+
     """
 
     return conversion_factor(init_units, final_units, energy_units_names, energy_units_vals)

src/perturbopy/postproc/utils/memory.py

@@ -0,0 +1,60 @@
+"""
+Memory management utilities
+"""
+
+import sys
+import numpy as np
+
+
+def get_size(array, name='array', dump=True):
+    """
+    Get size of an array in MB
+
+    Parameters
+    ----------
+    array : numpy array
+        Array.
+
+    name :
+        Name of array, optional.
+
+    dump :
+        Specifies whether to print out or not the size
+
+    Returns
+    -------
+    size : float
+        Size in MB.
+    """
+
+    # Check if the array is a numpy array and calculate size accordingly
+    if isinstance(array, np.ndarray):
+        # Calculate size in bytes for numpy array
+        size_bytes = array.size * array.itemsize
+    else:
+        # Use sys.getsizeof() for other types of arrays
+        size_bytes = sys.getsizeof(array)
+
+    size_kb = size_bytes / 1024.0
+
+    if size_kb < 1024.0:
+
+        size = size_kb
+        unit = 'KB'
+
+    elif size_kb / 1024.0 < 1024.0:
+
+        size_mb = size_kb / 1024.0
+        size = size_mb
+        unit = 'MB'
+
+    else:
+
+        size_gb = size_kb / 1024.0**2
+        size = size_gb
+        unit = 'GB'
+
+    if dump:
+        print(f'===Size of {name:<10} {str(array.shape):<12} {str(array.dtype):<8}: {size:6.3f} {unit}')
+
+    return size, unit

src/perturbopy/postproc/utils/plot_tools.py

@@ -76,15 +76,28 @@ def plot_recip_pt_labels(ax, labels, point_array, path_array, label_height="lowe
     elif label_height == "lower":
         label_height = ax.get_ylim()[0] - (ax.get_ylim()[1] - ax.get_ylim()[0]) * 0.1
 
+    # Create lists to store tick positions and labels
+    tick_positions = []
+    tick_labels = []
+
     for label in labels.keys():
         path_to_label = lattice.convert_point2path(labels[label], point_array, path_array)
 
         if path_to_label is None:
             continue
+
+        # Add positions and labels to lists
+        tick_positions.extend(path_to_label)
+        tick_labels.extend([label] * len(path_to_label))
+
         for x in path_to_label:
             if show_line:
                 ax.axvline(x, c="lightgray")
-                ax.text(x=x, y=label_height, s=label)
+                # ax.text(x=x, y=label_height, s=label)
+
+    # Set tick positions and labels on the x-axis
+    ax.set_xticks(tick_positions)
+    ax.set_xticklabels(tick_labels)
 
     return ax
 
@@ -221,7 +234,7 @@ def plot_vals_on_bands(ax, path, energies, energy_units, values, cmap='YlOrRd',
         vmin = min([min(values[key]) for key in values.keys()])
         vmax = max([max(values[key]) for key in values.keys()])
         norm = plt.Normalize(vmin, vmax)
-    
+
     for n in energies.keys():
 
         x = np.array(path)