Hey, we're making an EKG simulator but the QRS complex isn't handling the duration slider well.

When we increase the duration, we get a sharp increase, a level off around R, then a sharp decrease.

Can you take a look at the code for the complex and see if you can figure out why it's not a smooth hump, but so abrupt?

It's also not just the QRS complex, when we fiddle with some sliders, we get some jolting/jumping around of the EKG line, almost as if it's snapping to values

===

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Slider, CheckButtons, RadioButtons
from scipy.signal import butter, filtfilt
from scipy.special import erf  # Import error function for asymmetric Gaussian


def butter_bandpass(lowcut, highcut, fs, order=2):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a


class ECGParameters:
    def __init__(self):
        self.heart_rate = 75  # bpm
        self.rhythm_regular = True  # True for regular rhythm, False for irregular

        # P Wave Parameters
        self.p_wave_present = True
        self.p_wave_amp = 0.15
        self.p_wave_duration = 0.09  # seconds
        self.pr_interval = 0.16 # seconds

        # QRS Complex Parameters
        self.qrs_complex_present = True
        self.qrs_complex_amp = 1.0
        self.qrs_complex_duration = 0.09  # seconds

        # T Wave Parameters
        self.t_wave_present = True
        self.t_wave_amp = 0.3
        self.t_wave_duration = 0.16  # seconds
        self.t_wave_skew = 0.5  # Controls the asymmetry of the T wave

        # U Wave Parameters
        self.u_wave_present = False
        self.u_wave_amp = 0.1
        self.u_wave_duration = 0.04  # seconds

        # ST Segment Shift
        self.st_segment_shift = 0.0  # mV

        # Baseline Wander and Noise
        self.baseline_wander_amp = 0 #0.05
        self.noise_amp = 0 #0.02

        # Arrhythmia and Ectopic Beats
        self.fibrillation = False
        self.flutter = False
        self.ectopic_beats = False
        self.ectopic_beat_rate = 0.05
        self.arrhythmia = 'Normal Sinus Rhythm'


class ECGSimulator:
    def __init__(self, rate=75, duration=10, sampling_rate=500, seed=None):
        self.duration = duration
        self.sampling_rate = sampling_rate
        self.seed = seed
        self.t = []
        self.ecg_buffer_raw = []
        self.ecg_buffer_filtered = []
        self.b, self.a = butter_bandpass(0.5, 40.0, sampling_rate)
        self.current_time = 0.0
        self.dt = 1.0 / sampling_rate
        self.ecg_params = ECGParameters()  # Initialize parameters
        self.ecg_params.heart_rate = rate

        # Initialize beat timing
        self.rate = self.ecg_params.heart_rate
        self.beat_period = 60 / self.rate
        self.beat_start_time = 0.0
        self.next_beat_time = self.beat_start_time + self.beat_period
        self.beat_period_variation = self.beat_period

    def generate_asymmetric_gaussian(self, t, center, amp, width, skew):
        sigma = width / 2.355
        norm_factor = amp
        exponent = -((t - center) ** 2) / (2 * sigma ** 2)
        skew_factor = 1 + erf(skew * (t - center) / (sigma * np.sqrt(2)))
        return norm_factor * np.exp(exponent) * skew_factor

    def generate_p_wave(self, time_since_beat):
        params = self.ecg_params
        if not params.p_wave_present:
            return 0
        center = params.p_wave_center
        width = params.p_wave_duration
        amp = params.p_wave_amp
        sigma = width / 2.355
        return amp * np.exp(-((time_since_beat - center) ** 2) / (2 * sigma ** 2))

    def generate_qrs_complex(self, time_since_beat):
        params = self.ecg_params
        if not params.qrs_complex_present:
            return 0

        qrs_duration = params.qrs_complex_duration
        qrs_amp = params.qrs_complex_amp

        center = params.qrs_center

        # Q wave (make the width slightly larger for smoother transition)
        q_amp = -0.4 * qrs_amp  # Adjust for desired Q depth
        q_width = qrs_duration * 0.4  # Increase width to smooth out the spike
        q_center = center - qrs_duration * 0.25  # Place Q wave a bit earlier
        q_wave = q_amp * np.exp(-((time_since_beat - q_center) ** 2) / (2 * (q_width / 2.355) ** 2))

        # R wave (positive and broadest)
        r_amp = qrs_amp
        r_width = qrs_duration * 0.6  # Give R wave a broad base
        r_center = center  # R wave peaks in the middle
        r_wave = r_amp * np.exp(-((time_since_beat - r_center) ** 2) / (2 * (r_width / 2.355) ** 2))

        # S wave (increase width for smoother descent)
        s_amp = -0.4 * qrs_amp
        s_width = qrs_duration * 0.4  # Increase width to avoid sharp descent
        s_center = center + qrs_duration * 0.25  # Place S wave a bit later
        s_wave = s_amp * np.exp(-((time_since_beat - s_center) ** 2) / (2 * (s_width / 2.355) ** 2))

        # Only return the relevant wave component for the current time
        if time_since_beat < r_center:
            if time_since_beat < q_center:  # We're in the Q wave region
                return q_wave
            else:  # We're in the R wave region
                return r_wave
        else:  # We're in the S wave region
            return s_wave

    def generate_t_wave(self, time_since_beat):
        params = self.ecg_params
        if not params.t_wave_present:
            return 0
        center = params.t_wave_center
        width = params.t_wave_duration
        amp = params.t_wave_amp
        skew = params.t_wave_skew
        return self.generate_asymmetric_gaussian(time_since_beat, center, amp, width, skew)

    def generate_u_wave(self, time_since_beat):
        params = self.ecg_params
        if not params.u_wave_present:
            return 0
        center = params.u_wave_center
        width = params.u_wave_duration
        amp = params.u_wave_amp
        sigma = width / 2.355
        return amp * np.exp(-((time_since_beat - center) ** 2) / (2 * sigma ** 2))

    def generate_ecg_sample(self, t_current):
        params = self.ecg_params

        # Update heart rate and beat period
        self.rate = params.heart_rate
        self.beat_period = 60 / self.rate

        # Check if we have reached the next beat
        if t_current >= self.next_beat_time:
            self.beat_start_time = self.next_beat_time
            self.next_beat_time = self.beat_start_time + self.beat_period_variation

        # Time since last beat (relative to beat start)
        time_since_beat = t_current - self.beat_start_time

        # Adjust wave centers based on beat period
        params.p_wave_center = 0.1 * self.beat_period_variation
        params.qrs_center = params.p_wave_center + params.pr_interval
        params.t_wave_center = params.qrs_center + 0.2 * self.beat_period_variation
        params.u_wave_center = params.t_wave_center + 0.15 * self.beat_period_variation

        # Generate ECG components
        p_wave = self.generate_p_wave(time_since_beat)
        qrs_complex = self.generate_qrs_complex(time_since_beat)  # Already broken into Q, R, S
        t_wave = self.generate_t_wave(time_since_beat)
        u_wave = self.generate_u_wave(time_since_beat)

        # ST segment shift
        st_shift = params.st_segment_shift if params.qrs_center <= time_since_beat <= params.t_wave_center else 0

        # Add baseline wander and noise
        noise = params.noise_amp * np.random.randn()
        baseline_wander = params.baseline_wander_amp * np.sin(2 * np.pi * 0.33 * t_current)

        # Combine the ECG waveform by adding each component separately
        ecg_sample = p_wave + qrs_complex + t_wave + u_wave + baseline_wander + noise + st_shift

        return ecg_sample

    def apply_filter(self, ecg_data):
        return filtfilt(self.b, self.a, ecg_data)

    def init_widgets(self):
        # Adjust the positions of the widgets and make enough space at the bottom
        plt.subplots_adjust(left=0.15, bottom=0.6)  # Increase bottom margin further to make space for widgets

        # Heart Rate Slider
        widget_ax_slider_hr = plt.axes([0.05, 0.05, 0.4, 0.03], facecolor='lightgoldenrodyellow')  # Move to left
        self.slider_hr = Slider(widget_ax_slider_hr, 'Heart Rate', 30, 250, valinit=self.ecg_params.heart_rate,
                                valstep=1)

        # Wave Presence Checkboxes
        widget_ax_checkbox_waves = plt.axes([0.05, 0.7, 0.2, 0.2], facecolor='lightgoldenrodyellow')
        labels = ['P Wave', 'QRS Complex', 'T Wave', 'U Wave']
        actives = [self.ecg_params.p_wave_present, self.ecg_params.qrs_complex_present, self.ecg_params.t_wave_present,
                   self.ecg_params.u_wave_present]
        self.checkbox_waves = CheckButtons(widget_ax_checkbox_waves, labels, actives)

        # Amplitude Sliders (Column 1 - Left)
        amp_sliders_y = 0.35  # Y position for amplitude sliders
        amp_slider_height = 0.03  # Height for amplitude sliders
        self.slider_p_amp = Slider(plt.axes([0.05, amp_sliders_y, 0.4, amp_slider_height]), 'P Amp (mV)', 0.0, 1.0,
                                   valinit=self.ecg_params.p_wave_amp)
        self.slider_qrs_amp = Slider(plt.axes([0.05, amp_sliders_y - 0.06, 0.4, amp_slider_height]), 'QRS Amp (mV)',
                                     0.0, 2.0, valinit=self.ecg_params.qrs_complex_amp)
        self.slider_t_amp = Slider(plt.axes([0.05, amp_sliders_y - 0.12, 0.4, amp_slider_height]), 'T Amp (mV)', 0.0,
                                   1.0, valinit=self.ecg_params.t_wave_amp)
        self.slider_u_amp = Slider(plt.axes([0.05, amp_sliders_y - 0.18, 0.4, amp_slider_height]), 'U Amp (mV)', 0.0,
                                   0.5, valinit=self.ecg_params.u_wave_amp)

        # Duration Sliders (Column 2 - Right)
        dur_sliders_y = 0.35  # Y position for duration sliders
        dur_slider_height = 0.03  # Height for duration sliders
        self.slider_p_dur = Slider(plt.axes([0.55, dur_sliders_y, 0.4, dur_slider_height]), 'P Dur (s)', 0.04, 0.2,
                                   valinit=self.ecg_params.p_wave_duration)
        self.slider_qrs_dur = Slider(plt.axes([0.55, dur_sliders_y - 0.06, 0.4, dur_slider_height]), 'QRS Dur (s)',
                                     0.06, 1, valinit=self.ecg_params.qrs_complex_duration)
        self.slider_t_dur = Slider(plt.axes([0.55, dur_sliders_y - 0.12, 0.4, dur_slider_height]), 'T Dur (s)', 0.01,
                                   0.3, valinit=self.ecg_params.t_wave_duration)
        self.slider_u_dur = Slider(plt.axes([0.55, dur_sliders_y - 0.18, 0.4, dur_slider_height]), 'U Dur (s)', 0.02,
                                   0.1, valinit=self.ecg_params.u_wave_duration)

        # New Sliders for additional functionality (Column 2 - Right)
        self.slider_st_shift = Slider(plt.axes([0.55, dur_sliders_y - 0.24, 0.4, dur_slider_height]),
                                      'ST Segment Shift (mV)', -2.0, 2.0, valinit=self.ecg_params.st_segment_shift)
        self.slider_pr_interval = Slider(plt.axes([0.55, dur_sliders_y - 0.30, 0.4, dur_slider_height]),
                                         'P-R Interval (s)', 0.1, 0.3, valinit=self.ecg_params.pr_interval)
        self.slider_qt_interval = Slider(plt.axes([0.55, dur_sliders_y - 0.36, 0.4, dur_slider_height]),
                                         'QT Interval (s)', 0.3, 0.6, valinit=(
                        self.ecg_params.qrs_complex_duration + self.ecg_params.t_wave_duration))
        self.slider_hrv = Slider(plt.axes([0.55, dur_sliders_y - 0.42, 0.4, dur_slider_height]),
                                 'Heart Rate Variability (s)', 0.0, 0.1, valinit=0.05)

        # Connect widget events
        self.slider_hr.on_changed(self.on_slider_hr_changed)
        self.checkbox_waves.on_clicked(self.on_checkbox_waves_clicked)
        self.slider_p_amp.on_changed(self.on_slider_p_amp_changed)
        self.slider_qrs_amp.on_changed(self.on_slider_qrs_amp_changed)
        self.slider_t_amp.on_changed(self.on_slider_t_amp_changed)
        self.slider_u_amp.on_changed(self.on_slider_u_amp_changed)
        self.slider_p_dur.on_changed(self.on_slider_p_dur_changed)
        self.slider_qrs_dur.on_changed(self.on_slider_qrs_dur_changed)
        self.slider_t_dur.on_changed(self.on_slider_t_dur_changed)
        self.slider_u_dur.on_changed(self.on_slider_u_dur_changed)

        self.slider_st_shift.on_changed(self.on_slider_st_shift_changed)
        self.slider_pr_interval.on_changed(self.on_slider_pr_interval_changed)
        self.slider_qt_interval.on_changed(self.on_slider_qt_interval_changed)
        self.slider_hrv.on_changed(self.on_slider_hrv_changed)

    def on_slider_hr_changed(self, val):
        self.ecg_params.heart_rate = val

    def on_checkbox_waves_clicked(self, label):
        index = ['P Wave', 'QRS Complex', 'T Wave', 'U Wave'].index(label)
        status = self.checkbox_waves.get_status()[index]
        if label == 'P Wave':
            self.ecg_params.p_wave_present = status
        elif label == 'QRS Complex':
            self.ecg_params.qrs_complex_present = status
        elif label == 'T Wave':
            self.ecg_params.t_wave_present = status
        elif label == 'U Wave':
            self.ecg_params.u_wave_present = status

    def on_slider_p_amp_changed(self, val):
        self.ecg_params.p_wave_amp = val

    def on_slider_qrs_amp_changed(self, val):
        self.ecg_params.qrs_complex_amp = val

    def on_slider_t_amp_changed(self, val):
        self.ecg_params.t_wave_amp = val

    def on_slider_u_amp_changed(self, val):
        self.ecg_params.u_wave_amp = val

    def on_slider_p_dur_changed(self, val):
        self.ecg_params.p_wave_duration = val

    def on_slider_qrs_dur_changed(self, val):
        self.ecg_params.qrs_complex_duration = val

    def on_slider_t_dur_changed(self, val):
        self.ecg_params.t_wave_duration = val

    def on_slider_u_dur_changed(self, val):
        self.ecg_params.u_wave_duration = val

    def on_slider_st_shift_changed(self, val):
        self.ecg_params.st_segment_shift = val

    def on_slider_pr_interval_changed(self, val):
        self.ecg_params.pr_interval = val

    def on_slider_qt_interval_changed(self, val):
        self.ecg_params.qrs_complex_duration = val * 0.4  # Adjust QRS proportionally
        self.ecg_params.t_wave_duration = val * 0.6  # Adjust T-wave proportionally

    def on_slider_hrv_changed(self, val):
        self.ecg_params.heart_rate_variability = val

    def classify_arrhythmia(self):
        params = self.ecg_params
        heart_rate = params.heart_rate
        p_wave_present = params.p_wave_present
        qrs_present = params.qrs_complex_present
        qrs_duration = params.qrs_complex_duration
        qrs_amp = params.qrs_complex_amp

        arrhythmia = ''

        if qrs_present and p_wave_present:
            arrhythmia = "Normal Sinus Rhythm"
        elif not qrs_present:
            arrhythmia = 'Asystole'
        elif heart_rate > 100 and not p_wave_present and qrs_duration > 0.12:
            arrhythmia = 'Ventricular Tachycardia'
        elif heart_rate < 60 and p_wave_present and qrs_duration <= 0.12:
            arrhythmia = 'Sinus Bradycardia'
        elif heart_rate > 100 and p_wave_present and qrs_duration <= 0.12:
            arrhythmia = 'Sinus Tachycardia'
        else:
            arrhythmia = 'Unclassified'

        return arrhythmia

    def animate(self):
        fig, ax = plt.subplots(figsize=(10, 6))
        plt.subplots_adjust(left=0.25, bottom=0.2)  # Adjust margins to fit widgets below

        ax.set_facecolor('#fff')
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("ECG Signal (mV)")
        ax.set_ylim(-2, 2)
        ax.set_xlim(-10, 0)

        line, = ax.plot([], [], lw=1.5, color='black')
        small_box_interval = 0.04
        large_box_interval = 0.2
        ax.set_xticks(np.arange(-10, 0.2, large_box_interval))
        ax.set_xticks(np.arange(-10, 0.2, small_box_interval), minor=True)
        ax.grid(which='major', color='#ff70ad', linestyle='-', linewidth=0.8)
        ax.grid(which='minor', color='lightcoral', linestyle='--', linewidth=0.5)

        text_arrhythmia = ax.text(0.02, 0.90, '', transform=ax.transAxes)
        text_hr = ax.text(0.02, 0.95, '', transform=ax.transAxes)

        def init():
            line.set_data([], [])
            text_arrhythmia.set_text('')
            text_hr.set_text('')
            return line, text_arrhythmia, text_hr

        def update(frame):
            t_current = self.current_time
            ecg_sample = self.generate_ecg_sample(t_current)
            self.ecg_buffer_raw.append(ecg_sample)
            self.t.append(t_current)

            # Keep only the last 10 seconds of data
            while self.t and self.t[0] < self.current_time - 10:
                self.ecg_buffer_raw.pop(0)
                self.t.pop(0)

            t_shifted = [t - self.current_time for t in self.t]
            line.set_data(t_shifted, self.ecg_buffer_raw)

            arrhythmia = self.classify_arrhythmia()
            text_arrhythmia.set_text(f'Arrhythmia: {arrhythmia}')
            text_hr.set_text(f'Heart Rate: {self.ecg_params.heart_rate} BPM')

            self.current_time += self.dt
            return line, text_arrhythmia, text_hr

        self.init_widgets()

        ani = FuncAnimation(fig, update, init_func=init, frames=int(self.duration * self.sampling_rate),
                            interval=1000 / self.sampling_rate, blit=True)

        plt.show()


# Run the ECG simulation with widgets and arrhythmia detection
simulator = ECGSimulator(rate=75, duration=10, seed=42)
simulator.animate()