import numpy as np
import matplotlib.pyplot as plt

import up_the_ramp_sim as project_tools
import importlib

importlib.reload(project_tools)

<module 'up_the_ramp_sim' from '/home/koji/OneDrive/JHU/research/TSST/one_over_f/up_the_ramp_sim.py'>

image_dim = (256, 256)
frame_num = 6
num_amplifiers = 4
one_over_f_noise_std = 5.0
one_over_f_alpha_exponent = 1.2
one_over_f_knee_frequency = None

one_f_noise_frames = project_tools.generate_2d_readout_1f_noise_frames(image_dim, frame_num, num_amplifiers=num_amplifiers, one_over_f_noise_std=one_over_f_noise_std, one_over_f_alpha_exponent=one_over_f_alpha_exponent, one_over_f_knee_frequency=one_over_f_knee_frequency, frames_correlation=False)

Only 1/f noise

group_times = np.arange(0, frame_num * 10, 10)
group_images = one_f_noise_frames
final_rate_image_simple, final_intercept_image_simple = project_tools.fit_ramps_to_rate_image(group_images, group_times)

fig = plt.figure(figsize=(4, 4))
project_tools.plot_image_on_axis(fig.add_subplot(1, 1, 1), final_rate_image_simple, "Rate (e-/s)", limit_abs=True, cmap="bwr")

plt.hist(final_rate_image_simple.flatten(), bins=100)
plt.xlabel("Rate (e-/s)")
plt.ylabel("Counts")
plt.title("Histogram of Rate Image")
plt.show()

ramp_rate_fit_image = final_rate_image_simple * group_times[..., np.newaxis, np.newaxis] + final_intercept_image_simple
residual_image = group_images - ramp_rate_fit_image

(group_images - residual_image)[0,0,0], ramp_rate_fit_image[0,0,0]

(-0.20512627329238153, -0.20512627329238128)

fig, axes = plt.subplots(4, 6, figsize=(18, 12))
for i in range(6):
    project_tools.plot_image_on_axis(axes[0, i], group_images[i], f"Group {i+1}", limit_abs=True, cmap="bwr")
    project_tools.plot_image_on_axis(axes[1, i], ramp_rate_fit_image[i], f"rate fit", limit_abs=True, cmap="bwr")
    project_tools.plot_image_on_axis(axes[2, i], residual_image[i], f"Group {i+1} - rate fit (residual)", limit_abs=True, cmap="bwr")
    project_tools.plot_image_on_axis(axes[3, i], group_images[i] - residual_image[i], f"1/f - Residual", limit_abs=True, cmap="bwr")


plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.show()

pixel_y, pixel_x = 75, 160  # Example pixel coordinates for testing

pixel_ramp_values = group_images[:, pixel_y, pixel_x]
pixel_ramp_values_cleaned = (group_images - residual_image)[:, pixel_y, pixel_x]
# pixel_ramp_values_cleaned = (residuals_dark)[:, pixel_y, pixel_x]

fit_coeffs = np.polyfit(group_times, pixel_ramp_values, 1)
slope, intercept = fit_coeffs[0], fit_coeffs[1]

fit_coeffs_cleaned = np.polyfit(group_times, pixel_ramp_values_cleaned, 1)
slope_cleaned, intercept_cleaned = fit_coeffs_cleaned[0], fit_coeffs_cleaned[1]


# plot cleaned pixel ramp values side by side
plt.figure(figsize=(12, 6))

# Plot original pixel ramp values
ax1 = plt.subplot(1, 2, 1)
plt.scatter(group_times, pixel_ramp_values, label=f"Group Values for pixel ({pixel_y},{pixel_x})", color="blue")
plt.plot(group_times, intercept + slope * group_times, label=f"Linear Fit (Slope={slope:.2f} e-/s)", color="red", linestyle="--")
plt.xlabel("Effective Group Time (s)")
plt.ylabel("Signal (e-)")
plt.title(f"Original Up-the-Ramp Data and Fit for Pixel ({pixel_y},{pixel_x})")
plt.legend()
plt.grid(True)

# Plot cleaned pixel ramp values
ax2 = plt.subplot(1, 2, 2, sharey=ax1)
plt.scatter(group_times, pixel_ramp_values_cleaned, label=f"Ramp - Rate fit Values for pixel ({pixel_y},{pixel_x})", color="blue")
plt.plot(group_times, intercept_cleaned + slope_cleaned * group_times, label=f"Linear Fit (Slope={slope_cleaned:.2f} e-/s)", color="red", linestyle="--")
plt.xlabel("Effective Group Time (s)")
plt.title(f"Cleaned Up-the-Ramp Data and Fit for Pixel ({pixel_y},{pixel_x})")
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

ramp_cleaned = group_images - residual_image
final_rate_image_simple_cleaned, _ = project_tools.fit_ramps_to_rate_image(ramp_cleaned, group_times)

lim = 0.3
fig, axes = plt.subplots(1, 3, figsize=(12, 3))
project_tools.plot_image_on_axis(axes[0], final_rate_image_simple, "Rate (e-/s)", limit_abs=True, cmap="bwr", vmin=-lim, vmax=lim)
project_tools.plot_image_on_axis(axes[1], final_rate_image_simple_cleaned, "Rate (e-/s) cleaned", limit_abs=True, cmap="bwr")
project_tools.plot_image_on_axis(axes[2], final_rate_image_simple - final_rate_image_simple_cleaned, "Diff", limit_abs=True, cmap="bwr")

Adding White Noise

sigma_whitenoise = 5
group_images_with_noise = group_images.copy()
group_images_with_noise += np.random.normal(0, sigma_whitenoise, group_images_with_noise.shape)
final_rate_image_with_noise, final_intercept_image_with_noise = project_tools.fit_ramps_to_rate_image(group_images_with_noise, group_times)

fig = plt.figure(figsize=(4, 4))
project_tools.plot_image_on_axis(fig.add_subplot(1, 1, 1), final_rate_image_with_noise, "White noise Rate (e-/s)", limit_abs=True, cmap="bwr")

from scipy.ndimage import gaussian_filter1d

ramp_rate_fit_image = final_rate_image_with_noise * group_times[..., np.newaxis, np.newaxis] + final_intercept_image_with_noise
residuals = group_images_with_noise - ramp_rate_fit_image
residuals_readout = project_tools.frame2readout(residuals)
residuals_readout_filtered = gaussian_filter1d(residuals_readout, sigma=1, axis=-1)
residuals_filtered = project_tools.readout2frame(residuals_readout_filtered, image_dim)

limit_abs = np.max(np.abs(residuals_filtered))
fig, axes = plt.subplots(5, 6, figsize=(18, 15))
for i in range(6):
    project_tools.plot_image_on_axis(axes[0, i], group_images_with_noise[i], f"Group {i+1}", limit_abs=True, cmap="bwr", vmin=-limit_abs, vmax=limit_abs)
    project_tools.plot_image_on_axis(axes[1, i], ramp_rate_fit_image[i], f"Ramp rate fit", limit_abs=True, cmap="bwr", vmin=-limit_abs, vmax=limit_abs)
    project_tools.plot_image_on_axis(axes[2, i], residuals[i], f"Residual (Ramp - rate fit)", limit_abs=True, cmap="bwr", vmin=-limit_abs, vmax=limit_abs)
    project_tools.plot_image_on_axis(axes[3, i], residuals_filtered[i], f"Residual filtered", limit_abs=True, cmap="bwr", vmin=-limit_abs, vmax=limit_abs)
    project_tools.plot_image_on_axis(axes[4, i], one_f_noise_frames[i] - residuals_filtered[i], f"1/f - Residual filtered", limit_abs=True, cmap="bwr", vmin=-limit_abs, vmax=limit_abs)

plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.show()

pixel_y, pixel_x = 75, 160  # Example pixel coordinates for testing

pixel_ramp_values = group_images_with_noise[:, pixel_y, pixel_x]
pixel_ramp_values_cleaned = (group_images_with_noise - residuals_filtered)[:, pixel_y, pixel_x]

fit_coeffs = np.polyfit(group_times, pixel_ramp_values, 1)
slope, intercept = fit_coeffs[0], fit_coeffs[1]

fit_coeffs_cleaned = np.polyfit(group_times, pixel_ramp_values_cleaned, 1)
slope_cleaned, intercept_cleaned = fit_coeffs_cleaned[0], fit_coeffs_cleaned[1]


# plot cleaned pixel ramp values side by side
plt.figure(figsize=(12, 6))

# Plot original pixel ramp values
ax1 = plt.subplot(1, 2, 1)
plt.scatter(group_times, pixel_ramp_values, label=f"Group Values for pixel ({pixel_y},{pixel_x})", color="blue")
plt.plot(group_times, intercept + slope * group_times, label=f"Linear Fit (Slope={slope:.2f} e-/s)", color="red", linestyle="--")
plt.xlabel("Effective Group Time (s)")
plt.ylabel("Signal (e-)")
plt.title(f"Original Up-the-Ramp Data and Fit for Pixel ({pixel_y},{pixel_x})")
plt.legend()
plt.grid(True)

# Plot cleaned pixel ramp values
ax2 = plt.subplot(1, 2, 2, sharey=ax1)
plt.scatter(group_times, pixel_ramp_values_cleaned, label=f"Ramp - Rate fit Values for pixel ({pixel_y},{pixel_x})", color="blue")
plt.plot(group_times, intercept_cleaned + slope_cleaned * group_times, label=f"Linear Fit (Slope={slope_cleaned:.2f} e-/s)", color="red", linestyle="--")
plt.xlabel("Effective Group Time (s)")
plt.title(f"Cleaned Up-the-Ramp Data and Fit for Pixel ({pixel_y},{pixel_x})")
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

final_rate_image_with_noise_cleaned, _ = project_tools.fit_ramps_to_rate_image(group_images_with_noise - residuals_filtered, group_times)

lim = 0.4
fig, axes = plt.subplots(1, 3, figsize=(12, 3))
project_tools.plot_image_on_axis(axes[0], final_rate_image_with_noise, "Rate (e-/s)", limit_abs=True, cmap="bwr", vmin=-lim, vmax=lim)
project_tools.plot_image_on_axis(axes[1], final_rate_image_with_noise_cleaned, "Rate (e-/s) cleaned", limit_abs=True, cmap="bwr")
project_tools.plot_image_on_axis(axes[2], final_rate_image_with_noise - final_rate_image_with_noise_cleaned, "Diff", limit_abs=True, cmap="bwr")