Refactor: More intermediate text files, highlights, useful plots

- A subset of modes are identified by index in the config file as
  "highlights". This spectrum is broadened separately from the total
  sum.

- reference.pdf now just compares fundamentals from both methods: this
  is a more useful sanity-check

- plot.pdf now plots highlighted modes against fundamental and ref
  spectra from Abins. This seems the most research-relevant output.

- Various minor improvements: more control over binning and plot
  ranges, plotting stylesheet, simpler scripts (but more of them)

config/config.yaml

@@ -10,7 +10,28 @@ abins:
     Instrument: "TOSCA"
     Setting: "All detectors (TOSCA)"
 
-plot:
-    intensity_scale: 0.5  # Arbitrary scale factor
+binning:
     x_min: 0              # Minimum frequency in cm-1
     x_max: 4000           # Maximum frequency in cm-1
+    n_bins: 4001
+
+plot:
+    x_min: 0                      # Minimum frequency in cm-1
+    x_max: 4000                   # Maximum frequency in cm-1
+    intensity_scale: 0.5          # Scale factor for mode sums (e.g. averaging factor over detector banks)
+    style: "tableau-colorblind10" # Matplotlib stylesheet
+
+highlights:  # 0-based indices of modes to highlight in plot
+    - 2      # (For ethanol example these are A' modes)
+    - 4      # COUNTING FROM 0 MAY NOT BE THE SAME CONVENTION AS YOUR INPUT FILE
+    - 5
+    - 6
+    - 8
+    - 10
+    - 11
+    - 13
+    - 14
+    - 15
+    - 17
+    - 18
+    - 20

workflow/Snakefile

@@ -3,9 +3,10 @@ configfile: "config/config.yaml"
 
 rule all:
     input:
-        "results/modes.db",
-        "results/plot.pdf",
+        "results/intensities_binned.txt",
+        "results/highlights_binned.txt",
         "results/reference.pdf",
+        "results/plot.pdf",
 
 
 rule get_all_intensities:
@@ -42,16 +43,60 @@ rule get_reference:
         "scripts/get_reference.py"
 
 
-rule plot_modes:
+rule write_modes:
     input:
         "results/modes.db",
     output:
-        plot="results/plot.pdf",
         csv="results/intensities.csv",
         txt="results/intensities.txt",
     params:
         plot=config["plot"],
+        selection=None,
+    conda:
+        "envs/mantid.yml"
+    script:
+        "scripts/write_modes.py"
+
+
+rule write_highlights:
+    input:
+        "results/modes.db",
+    output:
+        csv="results/highlights.csv",
+        txt="results/highlights.txt",
+    params:
+        plot=config["plot"],
+        selection=config["highlights"],
+    conda:
+        "envs/mantid.yml"
+    script:
+        "scripts/write_modes.py"
+
+
+rule bin_and_broaden:
+    """Bin modes to grid and apply resolution convolution"""
+    input:
+        "results/{prefix}.csv",
+    output:
+        "results/{prefix}_binned.txt",
+    params:
         instrument=config["abins"]["Instrument"],
+        bin_config=config["binning"],
+    conda:
+        "envs/mantid.yml"
+    script:
+        "scripts/bin_and_broaden.py"
+
+
+rule plot_modes:
+    input:
+        highlights="results/highlights_binned.txt",
+        fundamentals="results/ref_fundamentals.csv",
+        multiphonon="results/ref_multiphonon.csv",
+    output:
+        plot="results/plot.pdf",
+    params:
+        plot=config["plot"],
     conda:
         "envs/mantid.yml"
     script:
@@ -60,8 +105,8 @@ rule plot_modes:
 
 rule plot_reference:
     input:
+        summed="results/intensities_binned.txt",
         fundamentals="results/ref_fundamentals.csv",
-        multiphonon="results/ref_multiphonon.csv",
     output:
         plot="results/reference.pdf",
     params:

workflow/scripts/bin_and_broaden.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+from euphonic import ureg, Quantity, Spectrum1D
+import numpy as np
+from numpy.polynomial import Polynomial
+from snakemake.script import snakemake
+
+bin_config = snakemake.params.bin_config
+
+
+def tosca_broadening(x: Quantity) -> Quantity:
+    """Gaussian width parameter for TOSCA as a function of energy"""
+
+    poly = Polynomial([2.5, 5e-3, 1e-7])
+    return poly(x.to("1/cm").magnitude) * ureg("1/cm")
+
+
+def get_spectrum(x: list[float], y: list[float]) -> Spectrum1D:
+    """Get a binned and broadened spectrum from scattered x/y input"""
+    bins = np.linspace(bin_config["x_min"], bin_config["x_max"], bin_config["n_bins"])
+    bin_width = bins[1] - bins[0]
+
+    hist, _ = np.histogram(x, bins=bins, weights=y, density=False)
+    hist /= bin_width  # Correct intensity scaling for bins
+
+    if snakemake.params.instrument != "TOSCA":
+        raise Exception("Currently only TOSCA resolution function is available.")
+
+    spectrum = Spectrum1D(ureg.Quantity(bins, "1/cm"), ureg.Quantity(hist, "cm"))
+    spectrum = spectrum.broaden(
+        x_width=tosca_broadening, shape="gauss", width_convention="std"
+    )
+    return spectrum
+
+
+def read_csv(input_file: Path) -> tuple[np.ndarray, np.ndarray]:
+    return np.loadtxt(
+        input_file, delimiter=",", skiprows=1, usecols=(1, 2), unpack=True
+    )
+
+
+spectrum = get_spectrum(*read_csv(Path(snakemake.input[0])))
+spectrum.to_text_file(snakemake.output[0])

workflow/scripts/plot_modes.py

@@ -1,11 +1,5 @@
-from itertools import groupby
-from operator import itemgetter
 from pathlib import Path
 
-import dataset
-from euphonic import ureg, Quantity, Spectrum1D
-from euphonic.plot import plot_1d
-
 import matplotlib
 import numpy as np
 from snakemake.script import snakemake
@@ -14,79 +8,39 @@
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
+fig, ax = plt.subplots()
 
 plot_config = snakemake.params.plot
-db_file = Path(snakemake.input[0])
-
-
-def tosca_broadening(x: Quantity) -> Quantity:
-    """Gaussian width parameter for TOSCA as a function of energy"""
-    from numpy.polynomial import Polynomial
-
-    poly = Polynomial([2.5, 5e-3, 1e-7])
-    return poly(x.to("1/cm").magnitude) * ureg("1/cm")
-
-
-def plot_spectrum(x: list[float], y: list[float]) -> None:
-    """Plot a binned spectrum line from scattered x/y input"""
-    bins = np.linspace(plot_config["x_min"], plot_config["x_max"], 4000)
-    bin_width = bins[1] - bins[0]
-
-    hist, _ = np.histogram(x, bins=bins, weights=y, density=False)
-    hist /= bin_width  # Correct intensity scaling for bins
-
-    if snakemake.params.instrument != "TOSCA":
-        raise Exception("Currently only TOSCA resolution function is available.")
-
-    spectrum = Spectrum1D(ureg.Quantity(bins, "1/cm"), ureg.Quantity(hist, "cm"))
-    spectrum = spectrum.broaden(
-        x_width=tosca_broadening, shape="gauss", width_convention="std"
-    )
-    spectrum *= plot_config["intensity_scale"]
-
-    fig = plot_1d(spectrum)
-    fig.gca().set_title("From individual mode sums")
-    fig.gca().set_xlim(bins[0], bins[-1])
-    fig.savefig(snakemake.output["plot"])
-
-
-# Create output directories if necessary
-for output in snakemake.output:
-    Path(output).parent.mkdir(parents=True, exist_ok=True)
-
-x, y = [], []
+plt.style.use(plot_config["style"])
 
-with (
-    dataset.connect(f"sqlite:///{db_file.resolve()}") as db,
-    open(snakemake.output["csv"], "wt") as csv,
-    open(snakemake.output["txt"], "wt") as txt,
-):
-    print(f"# mode_index,frequency,intensity", file=csv)
-    print(f"# mode_index  frequency  intensity", file=txt)
 
-    abins_table = db["abins"]
-    atoms_table = db["atoms"]
-    modes_table = db["modes"]
+def read_spectrum(input_file: Path, **kwargs) -> tuple[np.ndarray, np.ndarray]:
+    return np.loadtxt(input_file, unpack=True, usecols=(0, 1), **kwargs)
 
-    cross_sections = dict(
-        map(itemgetter("atom_index", "cross_section"), atoms_table.all())
-    )
 
-    frequencies = dict(map(itemgetter("mode_index", "frequency"), modes_table.all()))
+fig, ax = plt.subplots()
 
-    def get_intensity(data_row):
-        return cross_sections[data_row["atom_index"]] * data_row["weight"]
+# Plot multiphonon spectrum from Abins
+ax.plot(
+    *read_spectrum(snakemake.input["multiphonon"], delimiter=",", skiprows=2),
+    label="Fundamentals + multiphonon"
+)
 
-    for mode_index, group in groupby(
-        abins_table.find(order_by="mode_index"), key=itemgetter("mode_index")
-    ):
-        frequency = frequencies[mode_index]
-        intensity = sum(map(get_intensity, group))
+# Plot fundamental spectrum summed by Abins
+ax.plot(
+    *read_spectrum(snakemake.input["fundamentals"], delimiter=",", skiprows=2),
+    label="Fundamentals"
+)
 
-        print(f"{mode_index:d},{frequency:f},{intensity:f}", file=csv)
-        print(f"{mode_index:4d} {frequency:8.2f} {intensity:7.3f}", file=txt)
+# Plot contribution from selected modes
+x, y = read_spectrum(snakemake.input["highlights"], delimiter=" ")
+ax.plot(x, y * plot_config["intensity_scale"], label="Selected modes")
 
-        x.append(frequency)
-        y.append(intensity)
 
-plot_spectrum(x, y)
+ax.set_title("Simulated INS spectrum")
+ax.legend()
+ax.set_xlabel("Frequency / cm$^{-1}$")
+ax.set_xlim(plot_config["x_min"], plot_config["x_max"])
+ax.set_ylabel("Intensity")
+fig.tight_layout()
+fig.savefig(snakemake.output["plot"])

workflow/scripts/plot_reference.py

@@ -9,21 +9,26 @@
 
 
 plot_config = snakemake.params.plot
+plt.style.use(plot_config["style"])
 
 
-def read_spectrum(input_file: Path) -> tuple[np.ndarray, np.ndarray]:
-    return np.loadtxt(
-        input_file, delimiter=",", skiprows=2, usecols=(0, 1), unpack=True
-    )
+def read_spectrum(input_file: Path, **kwargs) -> tuple[np.ndarray, np.ndarray]:
+    return np.loadtxt(input_file, unpack=True, usecols=(0, 1), **kwargs)
 
 
 fig, ax = plt.subplots()
 
-ax.plot(*read_spectrum(snakemake.input["multiphonon"]), label="Multi-phonon")
-ax.plot(*read_spectrum(snakemake.input["fundamentals"]), label="Fundamentals")
+x, y = read_spectrum(snakemake.input["summed"], delimiter=" ")
+ax.plot(x, np.array(y) * plot_config["intensity_scale"], label="Sum over modes")
+
+ax.plot(
+    *read_spectrum(snakemake.input["fundamentals"], skiprows=2, delimiter=","),
+    label="From Abins",
+    linestyle=":",
+)
 
 ax.set_xlim(plot_config["x_min"], plot_config["x_max"])
-ax.set_title("Abins total spectrum")
+ax.set_title("Fundamental-only spectrum")
 ax.legend()
 
 fig.tight_layout()

workflow/scripts/write_modes.py

@@ -0,0 +1,43 @@
+from itertools import groupby
+from operator import itemgetter
+from pathlib import Path
+
+import dataset
+from snakemake.script import snakemake
+
+db_file = Path(snakemake.input[0])
+mode_selection = snakemake.params.selection
+
+with (
+    dataset.connect(f"sqlite:///{db_file.resolve()}") as db,
+    open(snakemake.output["csv"], "wt") as csv,
+    open(snakemake.output["txt"], "wt") as txt,
+):
+    print(f"# mode_index,frequency,intensity", file=csv)
+    print(f"# mode_index  frequency  intensity", file=txt)
+
+    abins_table = db["abins"]
+    atoms_table = db["atoms"]
+    modes_table = db["modes"]
+
+    cross_sections = dict(
+        map(itemgetter("atom_index", "cross_section"), atoms_table.all())
+    )
+
+    frequencies = dict(map(itemgetter("mode_index", "frequency"), modes_table.all()))
+
+    def get_intensity(data_row):
+        return cross_sections[data_row["atom_index"]] * data_row["weight"]
+
+    for mode_index, group in groupby(
+        abins_table.find(order_by="mode_index"), key=itemgetter("mode_index")
+    ):
+        if mode_selection is not None and mode_index not in mode_selection:
+            # Limit output to selected modes
+            continue
+
+        frequency = frequencies[mode_index]
+        intensity = sum(map(get_intensity, group))
+
+        print(f"{mode_index:d},{frequency:f},{intensity:f}", file=csv)
+        print(f"{mode_index:4d} {frequency:8.2f} {intensity:7.3f}", file=txt)