GeNoise_upd.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.ar_model import AutoReg
from joblib import Parallel, delayed

def generate_ar_noise(model_params, initial_noise, length, noise_std):
    noise = np.zeros(length)
    noise[:len(initial_noise)] = initial_noise

    # Vectorized computation for each step after initial noise
    for i in range(len(initial_noise), length):
        noise[i] = model_params[0] + np.dot(model_params[1:], noise[i - len(model_params[1:]):i][::-1])
        noise[i] += np.random.normal(scale=noise_std)  # Adding random noise

    return noise

def get_random_initial_noise(original_noise, lags):
    start_idx = np.random.randint(0, len(original_noise) - lags)
    return original_noise[start_idx:start_idx + lags]

def noise(noise_sample_length=1000, num_transients=10, plot=False):
    # Load waveform and prepare noise segments
    waveform = np.genfromtxt('waveform_data.txt', delimiter='\n') + 6.5
    noise_segments_indices = [(0, 3000), (4000, 6000)]
    noise_segments = [waveform[start:end] for start, end in noise_segments_indices]
    all_noise = np.concatenate(noise_segments)

    noise_std = np.std(all_noise) / 2
    lags = np.random.randint(1, 30, 1)[0]
    
    # Fit AR model
    ar_model = AutoReg(all_noise, lags=lags).fit()
    model_params = ar_model.params

    # Generate noise in parallel for each transient
    noises = Parallel(n_jobs=-1)(
        delayed(generate_ar_noise)(
            model_params,
            get_random_initial_noise(all_noise, lags),
            noise_sample_length,
            noise_std
        ) for _ in range(num_transients)
    )

    # Convert list to array
    noises = np.array(noises)

    if plot:
        plt.figure(figsize=(10, 4))
        plt.plot(noises[0])  # Plotting only the first transient
        plt.title('Generated Noise Using AR Model')
        plt.savefig("./noise.png")
    
    return noises