Fixed the INS q_planes plots and download

src/aiidalab_qe_vibroscopy/app/widgets/euphonicmodel.py

@@ -8,6 +8,8 @@
     produce_bands_weigthed_data,
     produce_powder_data,
     generated_curated_data,
+    produce_Q_section_spectrum,
+    produce_Q_section_modes,
 )
 
 from aiidalab_qe_vibroscopy.utils.euphonic.data.phonopy_interface import (
@@ -24,12 +26,6 @@
 )
 
 
-from aiidalab_qe_vibroscopy.utils.euphonic.tab_widgets.euphonic_q_planes_widgets import (
-    produce_Q_section_modes,
-    produce_Q_section_spectrum,
-)
-
-
 from aiidalab_qe.common.mvc import Model
 
 
@@ -360,13 +356,27 @@ def prepare_data_for_download(self):
         random_number = np.random.randint(0, 100)
 
         # Plotted_data
-        df = pd.DataFrame(self.z, index=self.y, columns=self.x)
+        if self.spectrum_type == "q_planes":
+            # we store x,y,z as values, not as values, indexes and columns
+            z = self.z.reshape(int(self.k_vec[-2]) + 1, int(self.h_vec[-2]) + 1)
+            df = pd.DataFrame(
+                z,
+                index=self.y[: int(self.k_vec[-2]) + 1],
+                columns=self.x[:: int(self.h_vec[-2]) + 1],
+            )
+        else:
+            df = pd.DataFrame(self.z, index=self.y, columns=self.x)
+
         data = base64.b64encode(df.to_csv().encode()).decode()
         filename = f"INS_structure_factor_{random_number}.csv"
 
         # model_state for template jinja plot script
         model_state = self.get_model_state()
         model_state["ylabel"] = self.ylabel
+        if hasattr(self, "xlabel"):
+            model_state["xlabel"] = self.xlabel
+        if hasattr(self, "Q0_vec"):
+            model_state["Q0"] = self.Q0_vec
         model_state["spectrum_type"] = self.spectrum_type
         if hasattr(self, "ticks_labels"):
             model_state["ticks_positions"] = self.ticks_positions

src/aiidalab_qe_vibroscopy/utils/euphonic/data/structure_factors.py

@@ -577,6 +577,111 @@ def update_max(max_val):
     matplotlib_save_or_show(save_filename=args.save_to)
 
 
+def produce_Q_section_modes(
+    fc,
+    h,
+    k,
+    Q0=np.array([0, 0, 0]),
+    n_h=100,
+    n_k=100,
+    h_extension=1,
+    k_extension=1,
+    temperature=0,
+):
+    from euphonic import ureg
+
+    # see get_Q_section
+    # h: array vector
+    # k: array vector
+    # Q0: "point" in Q-space used to build the portion of plane, using also the two vectors h and k.
+    # n_h, n_k: number of points along the two directions. or better, the two vectors.
+
+    def get_Q_section(h, k, Q0, n_h, n_k, h_extension, k_extension):
+        # every point in the space is Q=Q0+dv1*h+dv2*k, which
+
+        q_list = []
+        h_list = []
+        k_list = []
+
+        for dv1 in np.linspace(-h_extension, h_extension, n_h):
+            for dv2 in np.linspace(-k_extension, k_extension, n_k):
+                p = Q0 + dv1 * h + dv2 * k
+                q_list.append(p)  # Q list
+                h_list.append(dv1)  # *h[0])
+                k_list.append(dv2)  # *k[1])
+
+        return np.array(q_list), np.array(h_list), np.array(k_list)
+
+    q_array, h_array, k_array = get_Q_section(
+        h, k, Q0, n_h + 1, n_k + 1, h_extension, k_extension
+    )
+
+    modes = fc.calculate_qpoint_phonon_modes(qpts=q_array, reduce_qpts=False)
+
+    if temperature > 0:
+        blockPrint()
+        dw = _get_debye_waller(
+            temperature * ureg("K"),
+            fc,
+            # grid_spacing=(args.grid_spacing * recip_length_unit),
+            # **calc_modes_kwargs,
+        )
+        enablePrint()
+    else:
+        dw = None
+
+    labels = {
+        "q": f"Q0={[np.round(i,3) for i in Q0]}",
+        "h": f"h={[np.round(i,3) for i in h]}",
+        "k": f"k={[np.round(i,3) for i in k]}",
+    }
+
+    return modes, q_array, h_array, k_array, labels, dw
+
+
+def produce_Q_section_spectrum(
+    modes,
+    q_array,
+    h_array,
+    k_array,
+    ecenter,
+    deltaE=0.5,
+    bins=10,
+    spectrum_type="coherent",
+    dw=None,
+    labels=None,
+):
+    from aiidalab_qe_vibroscopy.utils.euphonic.data.structure_factors import (
+        blockPrint,
+        enablePrint,
+    )
+
+    # bins = 10 # hard coded beacuse here it does not change anything.
+    ebins = _get_energy_bins(
+        modes, bins + 1, emin=ecenter - deltaE, emax=ecenter + deltaE
+    )
+
+    blockPrint()
+    if (
+        spectrum_type == "coherent"
+    ):  # Temperature?? For now let's drop it otherwise it is complicated.
+        spectrum = modes.calculate_structure_factor(dw=dw).calculate_sqw_map(ebins)
+    elif spectrum_type == "dos":
+        spectrum = modes.calculate_dos_map(ebins)
+
+    mu = ecenter
+    sigma = (deltaE) / 2
+
+    # Gaussian weights.
+    weights = np.exp(-((spectrum.y_data.magnitude - mu) ** 2) / 2 * sigma**2) / np.sqrt(
+        2 * np.pi * sigma**2
+    )
+    av_spec = np.average(spectrum.z_data.magnitude, axis=1, weights=weights[:-1])
+    enablePrint()
+
+    return av_spec, q_array, h_array, k_array, labels
+
+
 def generated_curated_data(spectra):
     # here we concatenate the bands groups and create the ticks and labels.
 

src/aiidalab_qe_vibroscopy/utils/euphonic/plotting/plot_heatmap.py.j2

@@ -2,6 +2,7 @@ try:
     import pandas as pd
     import json
     import matplotlib.pyplot as plt
+    from matplotlib.ticker import MaxNLocator
 except:
     raise ImportError("Please install pandas, json and matplotlib to run this script.")
 import os
@@ -12,6 +13,7 @@ import os
 ## you can modify some of these: labels, cmap, filename. The rest is just a summary of the
 ## parameters used to produce the data.
 settings = {{ model_state }}
+spectrum_type = settings.get('spectrum_type', 'single_crystal')
 
 ## Load the heatmap data from a CSV file and plotting.
 # File paths
@@ -26,12 +28,23 @@ df = pd.read_csv(csv_file, index_col=0)
 
 # Create the heatmap
 plt.figure(figsize=(10, 8))
-plt.imshow(df.values, cmap=settings['cmap'], aspect='auto')
+
+plt.pcolormesh(df.columns, df.index, df.values, cmap=settings['cmap'], shading='auto')
+
+# Limit the number of ticks in the x axis
+## This is to avoid having too many ticks in the x-axis, which can make the plot unreadable
+## NB: you can change it at you convenience
+max_xticks = 9
+plt.gca().xaxis.set_major_locator(MaxNLocator(nbins=max_xticks))
 
 # Add y-axis label
 if 'ylabel' in settings:
     plt.ylabel(settings['ylabel'])
 
+# Add x-axis label
+if 'xlabel' in settings:
+    plt.xlabel(settings['xlabel'])
+
 # Add color bar
 plt.colorbar()
 
@@ -40,13 +53,13 @@ if 'ticks_positions' in settings and 'ticks_labels' in settings:
     plt.xticks(settings['ticks_positions'], settings['ticks_labels'])
 
 # Add title
-if 'spectrum_type' in settings:
-    plt.title(f"Inelastic neutron scattering data - {settings['spectrum_type']}")
+if "Q0" in settings:
+    plt.title(f"Inelastic neutron scattering data - {spectrum_type} - Q0 = {settings['Q0']}")
 else:
-    plt.title("Inelastic neutron scattering data")
+    plt.title(f"Inelastic neutron scattering data - {spectrum_type}")
 
 # Invert the y-axis
-plt.gca().invert_yaxis()
+if spectrum_type != "q_planes": plt.gca().invert_yaxis()
 
 # Show the plot
 plt.tight_layout()