scripts/basic_plotting.py

"""https://chatgpt.com/share/210d2fee-ca64-45a5-866e-e6df6e56bd1c"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle


# Function to plot spectra
def plot_spectra(sensor_data, num_points=100):
    wavelengths = np.array([410, 440, 470, 510, 550, 583, 620, 670])
    intensities = np.array(
        [
            sensor_data["ch410"],
            sensor_data["ch440"],
            sensor_data["ch470"],
            sensor_data["ch510"],
            sensor_data["ch550"],
            sensor_data["ch583"],
            sensor_data["ch620"],
            sensor_data["ch670"],
        ]
    )

    fig, ax = plt.subplots(figsize=(10, 6))

    # Adding rectangles for color bar effect
    dense_wavelengths = np.linspace(wavelengths.min(), wavelengths.max(), num_points)
    rect_height = max(intensities) * 0.02  # Height of the rectangles

    for dw in dense_wavelengths:
        rect = Rectangle(
            (
                dw - (wavelengths.max() - wavelengths.min()) / num_points / 2,
                -rect_height * 2,
            ),
            (wavelengths.max() - wavelengths.min()) / num_points,
            rect_height * 3,
            color=plt.cm.rainbow(
                (dw - wavelengths.min()) / (wavelengths.max() - wavelengths.min())
            ),
            edgecolor="none",
        )
        ax.add_patch(rect)

    # Main scatter plot
    scatter = ax.scatter(
        wavelengths, intensities, c=wavelengths, cmap="rainbow", edgecolor="k"
    )

    # Adding vertical lines from the x-axis to each point
    for wavelength, intensity in zip(wavelengths, intensities):
        ax.vlines(wavelength, 0, intensity, color="gray", linestyle="--", linewidth=1)

    ax.set_xlim(wavelengths.min() - 10, wavelengths.max() + 10)
    ax.set_ylim(0, max(intensities) + 10)  # Ensure the lower y limit is 0
    ax.set_xticks(wavelengths)
    ax.set_xlabel("Wavelength (nm)")
    ax.set_ylabel("Intensity")
    ax.set_title("Spectral Intensity vs. Wavelength")

    plt.show()


# Simulated sensor data
sensor_data = {
    "ch410": 10,
    "ch440": 20,
    "ch470": 15,
    "ch510": 30,
    "ch550": 25,
    "ch583": 40,
    "ch620": 35,
    "ch670": 50,
}

# Run the plot function with 100 points for the rectangles
plot_spectra(sensor_data, num_points=100)