simulate.py

#  %%
#!/usr/bin/env python
# coding: utf-8

from skimage.exposure import rescale_intensity
from PIL import Image, ImageOps
import pathlib
from dask.diagnostics import ProgressBar
import dask.array as da
import dask
from skimage.util import random_noise
import argparse
from sklearn import preprocessing
from scipy.sparse.linalg.interface import LinearOperator
from numpy.random import binomial
from sklearn.metrics import mean_squared_error, mean_absolute_error
from skimage.metrics import structural_similarity as ssim
import pandas as pd
from skimage.metrics import peak_signal_noise_ratio

import numpy as np
import matplotlib.pyplot as plt
import PIL
import scipy
from IPython.display import Video
from tqdm import tqdm

# from tqdm.notebook import tqdm

import multiprocessing
import os
import skimage
from enum import Enum, auto

from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer, SimpleImputer
from sklearn.metrics import mean_squared_error


# from sklearn import neural_network, metrics, gaussian_process, preprocessing,
#  svm, neighbors
# from sklearn import pipeline, model_selection

# from keras import metrics
# from keras import backend as K
from scipy.stats import pearsonr
from sklearn import svm, linear_model

# import microscPSF.microscPSF as msPSF
from skimage.metrics import structural_similarity as ssim

from scipy import matrix
from scipy.sparse import coo_matrix
import time
from scipy import linalg
from skimage import color, data, restoration, exposure
from skimage.transform import rescale, resize, downscale_local_mean
from scipy.signal import convolve2d as conv2

# import matlab.engine

# from sklearn.preprocessing import Imputer
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer

# from sklearn.experimental import enable_iterative_imputer
from scipy import signal
from sklearn.impute import SimpleImputer


class psf_switch_enum(Enum):
    STATIC, VAR_PSF, VAR_ILL = auto(), auto(), auto()


signal_strength = 1
coin_flip_bias = 0.5
savefig = 1
save_images = 0
image_scale = 4
psf_scale = 0.75
psf_gradient = 1
psf_size = 64
psf_type = "static"
max_iter = 100
thinning_type = "poisson"
out_dir = "out"
background_l = 1
background_k = 0

# Define constants: psf height width and image rescaling factor
# %%

# if __name__ == "__main__":

parser = argparse.ArgumentParser()
parser.add_argument("--signal_strength", default=signal_strength, type=float)
parser.add_argument("--coin_flip_bias", default=coin_flip_bias, type=float)
parser.add_argument("--savefig", default=savefig, type=int)
parser.add_argument("--save_images", default=save_images, type=int)
parser.add_argument("--image_scale", default=image_scale, type=int)
parser.add_argument("--psf_scale", default=psf_scale, type=float)
parser.add_argument("--psf_gradient", default=psf_gradient, type=float)
parser.add_argument("--psf_type", default=psf_type, type=str)
parser.add_argument("--max_iter", default=max_iter, type=int)
parser.add_argument("--thinning_type", default=thinning_type, type=str)
parser.add_argument("--out_dir", default=out_dir, type=str)
parser.add_argument("--background_l", default=background_l, type=float)
parser.add_argument("--background_k", default=background_k, type=float)

try:
    args = parser.parse_args()
# globals().update(vars(args))
except:
    args = parser.parse_args([])
globals().update(vars(args))
print(vars(args))
# if __name__ == "__main__":
#     args = parser.parse_args()
# %%
if psf_type == "static":
    psf_switch = psf_switch_enum.STATIC
elif psf_type == "variable":
    psf_switch = psf_switch_enum.VAR_PSF
else:
    raise Exception("I don't recongise this psf type")

# # Image formation

psf_w, psf_h, scale = psf_size, psf_size, image_scale
psf_window_w, psf_window_h = round(psf_w / scale), round(psf_h / scale)
# Define approximate PSF function


def gaussian(x, mu, sig):
    return np.exp(-np.power(x - mu, 2.0) / (2 * np.power(sig, 2.0)))


def psf_guass(w=psf_w, h=psf_h, sigma=3):
    # blank_psf = np.zeros((w,h))
    xx, yy = np.meshgrid(np.linspace(-1, 1, w), np.linspace(-1, 1, h))
    psf = gaussian(xx, 0, sigma) * gaussian(yy, 0, sigma)
    return psf / psf.sum()  # Normalise PSF "energy"


# Define a function that scales the PSF as a function of radial distance


def psf_vary(psf_window_h, psf_window_w, radius, psf_scale, psf_gradient, scale):
    sigma_0 = psf_scale / scale
    sigma = psf_gradient*abs(radius)*sigma_0 + sigma_0
    return psf_guass(
        w=round(psf_window_h),
        h=round(psf_window_w),
        sigma=sigma
    )


static_psf = psf_guass(
    w=round(psf_window_h), h=round(psf_window_w), sigma=psf_scale / scale
)
#  %%
# %% Deconvolution

astro = (
    rescale(color.rgb2gray(data.human_mitosis()),
            1.0 / image_scale)
)

# image_width, image_height = np.shape(astro)

# x_range = np.linspace(-1, 1, image_width)
# y_range = np.linspace(-1, 1, image_height)

# xx, yy = np.meshgrid(x_range, y_range)
# out = np.square(xx, yy)
# disk = np.sqrt(
#     np.add(
#         *np.square(
#             [
#                 *np.meshgrid(
#                     np.linspace(-1, 1, image_width), np.linspace(-1, 1, image_height)
#                 )
#             ]
#         )
#     )
# )

# #  %%

# xx, yy = np.meshgrid(
#                             np.linspace(-1, 1, image_width),
#                             np.linspace(-1, 1, image_height),
#                         )
# disk = (
#     np.sqrt(
#         np.sum(
#             [
#                 *np.square(
#                     [
#                         xx, yy
#                     ]
#                 )
#             ],
#             axis=0,
#         )
#     )
#     < psf_scale
# )

# plt.imshow(disk)

# # %%
# static_psf_full = np.fft.ifftshift(np.fft.fft2(disk))

# # np.absolute(xx) > psf_width

# # 128*((1/0.75)*(1/64))*2

# psf_width_pc = (((1/psf_scale)*(1/(image_width/2))))
# psf_width_px = 20*image_width*psf_width_pc

# # psf_mask = ~np.prod((np.absolute([xx,yy]) < 0.1),axis=0).astype(bool)

# # import numpy.ma as ma

# # mx = ma.masked_array(static_psf_full, mask=psf_mask)
# # ma.getdata(mx)
# # # static_psf_full[~psf_mask] = np.nan
# # # static_psf_full = static_psf_full[~np.isnan(static_psf_full)]

# plt.imshow(np.absolute(static_psf_full))
# # %%
# from skimage.util import crop
# # bounding_box = tuple(image_width+np.ceil(((-psf_width_px/2,psf_width_px/2))*len(static_psf_full.shape)));bounding_box

# # bounding_box = tuple(image_width+np.ceil(((-psf_width_px/2,psf_width_px/2))*len(static_psf_full.shape)));bounding_box

# # bounding_box = tuple(map(tuple,np.array([image_width+np.array([-psf_width_px/2,psf_width_px/2])]*len(static_psf_full.shape))))

# # bounding_box = np.array([image_width/2+np.array([-psf_width_px/2,psf_width_px/2])]*len(static_psf_full.shape)).astype(int)
# # lower = image_width/2-psf_width_px/2
# # upper = image_width/2+psf_width_px/2
# crop_edge = np.rint((image_width/2)-psf_width_px/2).astype(int)
# static_psf = crop(static_psf_full,((37,37),(37,37)))
# plt.imshow(np.real(static_psf))
# plt.title("PSF")
# plt.show()
# %%
astro_blur = conv2(astro, static_psf, "same")  # Blur image

# astro_noise = (
#     np.random.poisson(lam=25, size=astro_blur.shape)
# ) / 255.0  # Add Noise to Image


# astro_corrupt = astro_noisy = astro_blur + astro_noise  # Add Noise to Image

# astro_corrupt = (
#     np.random.poisson(lam=astro_blur * signal_strength, size=astro_blur.shape)
# ) / signal_strength

# astro_corrupt = random_noise(astro_blur, "poisson")

# deconvolved_RL = np.real(restoration.richardson_lucy(
#     astro_blur, static_psf, iterations=10
# ))  # RL deconvolution

# fig, ax = plt.subplots(nrows=1, ncols=4, figsize=(16, 7))
# ax[0].imshow(astro)
# ax[0].set_title("Truth")
# ax[1].imshow(astro_blur)
# ax[1].set_title("Blurred")
# ax[2].imshow(astro_corrupt)
# ax[2].set_title("Blurred and noised")
# ax[3].imshow(deconvolved_RL)
# ax[3].set_title("Deconvolved")
# plt.show()

# %%

N_v = np.ma.size(astro)
N_p = np.ma.size(astro_blur)
# measurement_matrix = matrix(np.zeros((N_p, N_v)))


zero_image = np.zeros_like(astro)
psf_window_volume = np.full((N_v, psf_window_w, psf_window_h), np.NaN)

x_astro, y_astro = astro_blur.shape
xx_astro, yy_astro = np.meshgrid(
    np.linspace(-1, 1, x_astro), np.linspace(-1, 1, y_astro)
)


# Make the PSF vary across the image (as a function of radius)

r_map = np.sqrt(xx_astro ** 2 + yy_astro ** 2)
radius_samples = np.linspace(-1, 1, 5)
fig, ax = plt.subplots(nrows=1, ncols=len(radius_samples), figsize=(16, 7))

for i, radius in enumerate(np.linspace(-1, 1, 5)):
    psf_current = psf_vary(psf_window_h, psf_window_w,
                           radius, psf_scale, psf_gradient, scale)
    ax[i].imshow(psf_current)
    ax[i].set_title("Radius: " + str(radius))
plt.show()


# if psf_switch == psf_switch_enum.STATIC:
#     filename = "H_staticpsf"
# if psf_switch == psf_switch_enum.VAR_PSF:
#     filename = "H_varpsf"

# %% Loop over each row of H and insert a flattened PSF
# that is the same shape as the input image


def make_flat_psf(i, psf_window_h, psf_window_w,
                  psf_scale, psf_gradient, scale, astro):
    coords = np.unravel_index(i, np.array(astro.shape))
    r_dist = r_map[coords]  # Convert to radius
    # Select mode for generating H, i.e. static/varying psf etc.
    # if psf_switch == psf_switch_enum.STATIC:
    #     psf_current = static_psf
    # if psf_switch == psf_switch_enum.VAR_PSF:
    psf_current = psf_vary(psf_window_h, psf_window_w,
                           r_dist, psf_scale, psf_gradient, scale)

    psf_window_volume[i, :, :] = psf_current
    delta_image = np.zeros_like(astro)
    delta_image[np.unravel_index(i, astro.shape)] = 1
    # Convolve PSF with a image with a single 1 at coord
    delta_PSF = scipy.ndimage.convolve(delta_image, psf_current)
    return delta_PSF.flatten()


delayed_list = []
for i in tqdm(np.arange(N_v)):
    # Get the xy coordinates of the ith pixel in the original image
    # delayed_list.append(
    delayed_obj = dask.delayed(make_flat_psf)(i, psf_window_h, psf_window_w,
                                              psf_scale, psf_gradient, scale, astro)
    delayed_list.append(
        da.from_delayed(
            delayed_obj, shape=(np.multiply(*astro.shape),), dtype=np.float32
        )
    )
    # )
    # measurement_matrix[i, :] = delta_PSF.flatten()
# array = da.from_delayed(
#     delayed_list, (len(delayed_list),), dtype=float
# )

with ProgressBar():
    stack = da.stack(delayed_list, axis=0)
    measurement_matrix = np.array(stack)
#  stack.compute()
# %%
# measurement_matrix = stack
plt.figure(figsize=(18, 7))
plt.imshow(exposure.equalize_hist(measurement_matrix), cmap="gray_r")
# %%

H = scipy.sparse.linalg.aslinearoperator(measurement_matrix)
# %%
g_blurred = H.dot(astro.reshape(-1, 1))
# f = np.matrix(g_blurred) + astro_noise.reshape(g_blurred.shape)
# g_blurred

lin = np.linspace(-1, 1, astro.shape[0])
xx, yy = np.meshgrid(lin, lin)

background = background_l/(1 + np.exp(-background_k*xx))
background = background - background.min()

g_blurred = g_blurred + background.reshape(g_blurred.shape)
plt.imshow(g_blurred.reshape(astro.shape))
#  %%
# Dodgy noise
f = random_noise(np.matrix(g_blurred), "poisson")
rng = np.random.default_rng()
# Better noise
f = rng.poisson(np.rint(g_blurred*(2**16)*signal_strength))/2**16
plt.imshow(f.reshape(astro.shape))
#  %%
fig, ax = plt.subplots(ncols=3, nrows=1, figsize=(16, 7))

ax[0].imshow(f.reshape(astro.shape))
ax[0].set_title("Blurred")
ax[1].imshow(astro)
ax[1].set_title("Original")
# ax[2].imshow(g.reshape(astro.shape));ax[2].set_title("RL")
plt.show()

#  %% Deconvolve matrices
A = H
b = f
x0 = None
Rtol = 1e-6
NE_Rtol = 1e-6
# max_iter = 100
sigmaSq = 0.0
beta = 0.0
# integer_signal = np.rint(b*2**16).astype(int)
coin_flip_scale = np.random.binomial(
    [2 ** 8] * len(b), coin_flip_bias) / [2 ** 8]
# coin_flip_scale = np.random.binomial([2 ** 8] * len(b), 0.5) / 2 ** 8


T_thinned = b.copy().reshape(astro.shape)
V_thinned = b.copy().reshape(astro.shape)

# thinning_type = "spatial_interlaced"

if thinning_type == "poisson":
    # %% FIsh thinning
    T_scaled = np.matrix(np.multiply(coin_flip_scale, np.array(b.copy()).T).T)
    V_scaled = b - T_scaled

if thinning_type == "spatial_interlaced":
    T_thinned = T_thinned.copy()
    V_thinned = V_thinned.copy()
    T_thinned[:, 0::2] = 0
    V_thinned[:, 1::2] = 0
    T_scaled = T_thinned
    V_scaled = V_thinned

if thinning_type == "spatial_interpolated":

    T_thinned_squished = T_thinned.copy()[:, 0::2]
    V_thinned_squished = V_thinned.copy()[:, 1::2]

    T_scaled = skimage.transform.rescale(
        T_thinned_squished, (1.0, 2.0), anti_aliasing=True
    )
    V_scaled = skimage.transform.rescale(
        V_thinned_squished, (1.0, 2.0), anti_aliasing=True
    )
    # T = T_scaled
    # V = V_scaled

if thinning_type == "spatial_repeated":
    T_thinned_squished = T_thinned.copy()
    V_thinned_squished = V_thinned.copy()

    T_thinned_squished[:, 0::2] = T_thinned_squished[:, 1::2]
    V_thinned_squished[:, 1::2] = V_thinned_squished[:, 0::2]

    T_scaled = T_thinned_squished
    V_scaled = V_thinned_squished

if thinning_type == "none":
    T_scaled = b.copy()
    V_scaled = b.copy()


T = T_scaled.reshape(b.shape)
V = V_scaled.reshape(b.shape)

# FIsh thinning
# T = np.matrix(np.multiply(coin_flip_scale, np.array(b).T).T)
# V = b - T

# T = np.matrix(np.multiply(coin_flip_scale, np.array(b).T).T)
# V = b - T

# Spatial thinning

# T_thinned = b.copy().reshape(astro.shape)
# V_thinned = b.copy().reshape(astro.shape)

# T_thinned_squished = T_thinned.copy()[:,0::2]
# V_thinned_squished = V_thinned.copy()[:,1::2]

# T_thinned_squished = T_thinned.copy()
# V_thinned_squished = V_thinned.copy()

# T_thinned_squished[:,0::2] = T_thinned_squished[:,1::2]
# V_thinned_squished[:,1::2] = V_thinned_squished[:,0::2]

# T_thinned[:,0::2] = 0
# V_thinned[:,1::2] = 0

# T_scaled = T_thinned_squished
# V_scaled = V_thinned_squished

b = np.matrix(T).reshape(f.shape)
gt = astro.reshape(V.shape)

plt.imshow(V.reshape(astro.shape))
plt.show()

plt.imshow(T.reshape(astro.shape))
plt.show()

plt.imshow(T.reshape(astro.shape))
plt.show()
# b = T

# The A operator represents a large, sparse matrix that
# has dimensions [ nrays x nvoxels ]

nrays = A.shape[0]
nvoxels = A.shape[1]

# Pre-compute some values for use in stopping criteria below
b_norm = np.linalg.norm(b)
trAb = A.rmatvec(b)
trAb_norm = np.linalg.norm(trAb)

# Start the optimization from the initial volume of a focal stack.
if x0 != None:
    x = x0
else:
    x = trAb

Rnrm = np.zeros(max_iter)
Rnrm_v = np.zeros(max_iter)
Rnrm_T = np.zeros(max_iter)
Rnrm_V = np.zeros(max_iter)

residual = np.zeros(max_iter)
residual_v = np.zeros(max_iter)
residual_V = np.zeros(max_iter)
residual_T = np.zeros(max_iter)

Xnrm = np.zeros(max_iter)
NE_Rnrm = np.zeros(max_iter)
gt_error_l1 = np.zeros(max_iter)
gt_error_l2 = np.zeros(max_iter)

V_error_l1 = np.zeros(max_iter)
T_error_l1 = np.zeros(max_iter)
cross_correlation = np.zeros(max_iter)
gt_error_ssim = np.zeros(max_iter)
log_liklihood = np.zeros(max_iter)
log_liklihood_v = np.zeros(max_iter)
log_liklihood_V = np.zeros(max_iter)
log_liklihood_T = np.zeros(max_iter)
psnr_V = np.zeros(max_iter)
psnr_T = np.zeros(max_iter)
psnr_b = np.zeros(max_iter)
kl_div = np.zeros(max_iter)

norm_v = np.zeros(max_iter)

eps = np.spacing(1)
tau = np.sqrt(eps)
sigsq = tau
minx = x.min()

# If initial guess has negative values, compensate
if minx < 0:
    x = x - min(0, minx) + sigsq

normalization = A.rmatvec(np.ones(nrays)) + 1
print(normalization.min(), normalization.max())
c = (
    A.matvec(x)
    + np.matrix(beta * np.ones(nrays)).T
    + np.matrix(sigmaSq * np.ones(nrays)).T
)
b = b + np.matrix(sigmaSq * np.ones(nrays)).T

# i =0

iterations = range(max_iter)
range_tqdm = tqdm(iterations)


def p_x_given_Ax(x, Ax):
    return np.multiply(np.log(Ax), (x)) - Ax - np.log(scipy.special.factorial(x))


def log_liklihood_x_given_Ax(x, Ax):
    return np.sum(
        p_x_given_Ax(x, Ax)
    )

#  %% RL


for i in range_tqdm:
    tic = time.time()
    x_prev = x

    # STEP 1: RL Update step
    # b / c
    # (c+1e-12).min()
    v = A.rmatvec(b / c)
    x = (np.multiply(x_prev, v)) / matrix(normalization).T
    if x.min() < 0:
        x = x - min(0, x.min()) + sigsq
    Ax = A.matvec(x)
    # b-Ax
    residual = np.linalg.norm(b - Ax)
    residual_v[i] = np.linalg.norm(v - Ax)
    residual_V[i] = np.linalg.norm(V - Ax)
    residual_T[i] = np.linalg.norm(T - Ax)
    norm_v[i] = np.linalg.norm(v)
    # residual_V = V - Ax

    c = (
        A.matvec(x)
        + np.matrix(beta * np.ones(nrays)).T
        + np.matrix(sigmaSq * np.ones(nrays)).T
    )
    # c = Ax + beta*np.ones(nrays) + sigmaSq*np.ones(nrays);

    # STEP 2: Compute residuals and check stopping criteria
    Rnrm[i] = residual / b_norm
    Xnrm[i] = np.linalg.norm(x - x_prev) / nvoxels
    Rnrm_V[i] = residual_V[i] / np.linalg.norm(V)
    Rnrm_T[i] = residual_T[i] / np.linalg.norm(T)
    Rnrm_v[i] = residual_v[i] / np.linalg.norm(v)

    # NE_Rnrm[i] = np.linalg.norm(trAb - A.rmatvec(Ax)) / trAb_norm
    # # disabled for now to save on extra rmatvec

    toc = time.time()
    range_tqdm.set_description_str(
        f"RL Iteration {i:0} ({(toc-tic):.4f} seconds)"
        f"Residual Norm: {Rnrm[i]:.4f}  (tol = {Rtol:.2f}"
        f"Residual V Norm: {Rnrm_V[i]:.5f} (tol = {Rtol:.2f}"
        f"Residual T Norm: {Rnrm_T[i]:.5f} (tol = {Rtol:.2f}"
        f"v Norm: {norm_v[i]:.4f} (tol = {Rtol:.2f}"
        # print '\tError Norm: %0.4g(tol = %0.2e)  ' % (NE_Rnrm[i], NE_Rtol)
        f"Update Norm:{Xnrm[i]:.4f}",
        refresh=False,
    )
    # range_tqdm.set_postfix({
    #     "RL Iteration":toc-tic
    # })
    # stop because residual satisfies ||b-A*x|| / ||b||<= Rtol
    if Rnrm[i] <= Rtol:
        break

    # stop because normal equations residual satisfies
    #  ||A'*b-A'*A*x|| / ||A'b||<= NE_Rtol
    # if NE_Rnrm[i] <= NE_Rtol:
    #    break
    gt_error_l2[i] = mean_squared_error(x, gt)
    gt_error_l1[i] = mean_absolute_error(x, gt)

    V_error_l1[i] = mean_absolute_error(Ax, V)
    T_error_l1[i] = mean_absolute_error(Ax, T)

    x_flat = np.array(x.copy()).flatten()
    gt_flat = np.array(gt.copy()).flatten()

    gt_error_ssim[i] = ssim(x_flat, gt_flat)

    x_flat_scaled = preprocessing.scale(x_flat)
    gt_flat_scaled = preprocessing.scale(gt_flat)

    # a_flat_scaled = (x_flat - np.mean(x_flat)) / (np.std(x_flat) * len(x_flat))
    # b_flat_scaled = (gt_flat - np.mean(gt_flat)) / (np.std(gt_flat))
    # cross_correlation[i] = np.sum(np.correlate(a_flat_scaled, b_flat_scaled, 'full'))
    cross_correlation[i] = np.sum(np.correlate(
        x_flat_scaled, gt_flat_scaled, "full"))
    # a = np.dot(abs(x_flat_scaled),abs(gt_flat_scaled),'full')

    log_liklihood[i] = log_liklihood_x_given_Ax(b, Ax)
    log_liklihood_v[i] = log_liklihood_x_given_Ax(v, Ax)
    log_liklihood_V[i] = log_liklihood_x_given_Ax(V, Ax)
    log_liklihood_T[i] = log_liklihood_x_given_Ax(T, Ax)
    psnr_V[i] = peak_signal_noise_ratio(V, Ax)
    psnr_T[i] = peak_signal_noise_ratio(T, Ax)
    psnr_b[i] = peak_signal_noise_ratio(b, Ax)
    kl_div[i] = np.sum(
        np.multiply(
            p_x_given_Ax(V, Ax),
            np.log(
                p_x_given_Ax(V, Ax) / p_x_given_Ax(T, Ax)
            )
        )
    )

# data = pd.DataFrame()
#  %%
vars_dict = dict(
    {
        i: eval(i)
        for i in [
            "iterations",
            "gt_error_l2",
            "gt_error_l1",
            "V_error_l1",
            "T_error_l1",
            "gt_error_ssim",
            "cross_correlation",
            "log_liklihood",
            "log_liklihood_v",
            "log_liklihood_V",
            "log_liklihood_T",
            "residual_v",
            "residual_V",
            "residual_T",
            "norm_v",
            "Rnrm",
            "Xnrm",
            "Rnrm_V",
            "Rnrm_T",
            "Rnrm_v",
            "psnr_V",
            "psnr_T",
            "psnr_b",
            "kl_div"
        ]
    }
)


print(f"dir: {out_dir}")
pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)

vars_dict.update(vars(args))
data_df = pd.DataFrame.from_dict(vars_dict)
# data.join(metadata)
# metadata

data_df.to_csv(os.path.join(out_dir, "data.csv"))


Image.fromarray(rescale_intensity(T.reshape(astro.shape), out_range=np.uint8)).convert(
    "L"
).save(os.path.join(out_dir, "T.png"))
Image.fromarray(rescale_intensity(f.reshape(astro.shape), out_range=np.uint8)).convert(
    "L"
).save(os.path.join(out_dir, "f.png"))
Image.fromarray(rescale_intensity(V.reshape(astro.shape), out_range=np.uint8)).convert(
    "L"
).save(os.path.join(out_dir, "V.png"))
Image.fromarray(
    rescale_intensity(g_blurred.reshape(astro.shape), out_range=np.uint8)
).convert("L").save(os.path.join(out_dir, "g_blurred.png"))
Image.fromarray(rescale_intensity(astro, out_range=np.uint8)).convert("L").save(
    os.path.join(out_dir, "raw.png")
)


#  %%
plt.plot(Rnrm_v[1:-1])
plt.yscale("log")
plt.title("v residuals")
plt.xlabel("Iterations")
plt.ylabel("Normalised residuals")
plt.show()

plt.plot(Rnrm_V[1:-1])
plt.yscale("log")
plt.title("V residuals")
plt.xlabel("Iterations")
plt.ylabel("Normalised residuals")
plt.show()

plt.plot(Rnrm_T[1:-1])
plt.yscale("log")
plt.title("T residuals")
plt.xlabel("Iterations")
plt.ylabel("Normalised residuals")
plt.show()

plt.plot(residual_v[1:-1])
plt.title("v residuals")
plt.xlabel("Iterations")
plt.ylabel(" residuals")
plt.show()

plt.plot(residual_V[1:-1])
plt.title("V residuals")
plt.xlabel("Iterations")
plt.ylabel(" residuals")
plt.show()

plt.plot(residual_T[1:-1])
plt.title("T residuals")
plt.xlabel("Iterations")
plt.ylabel(" residuals")
plt.show()

plt.plot(Rnrm[1:-1])
plt.title("Residuals")
plt.xlabel("Iterations")
plt.ylabel("Normalised residuals")
plt.show()

plt.plot(V_error_l1[1:-1])
plt.title("Ax v V ~ L2")
plt.xlabel("Iterations")
plt.ylabel("MSE")
plt.show()

plt.plot(V_error_l1[1:-1])
plt.title("Ax v T ~ L1")
plt.xlabel("Iterations")
plt.ylabel("MSE")
plt.show()

plt.plot(gt_error_l2[1:-1])
plt.title("GT v x ~ L1")
plt.xlabel("Iterations")
plt.ylabel("MSE")
plt.show()

plt.plot(gt_error_l1[1:-1])
plt.title("GT v x ~ L2")
plt.xlabel("Iterations")
plt.ylabel("MSE")
plt.show()

plt.plot(cross_correlation[1:-1])
plt.title("GT v x ~ cross_correlation")
plt.xlabel("Iterations")
plt.ylabel("MSE")
plt.show()


plt.plot(gt_error_ssim[1:-1])
plt.title("GT v x ~ SSIM")
plt.xlabel("Iterations")
plt.ylabel("SSIM")
plt.show()

plt.plot(log_liklihood[1:-1])
plt.title("GT v x ~ log_liklihood")
plt.xlabel("Iterations")
plt.ylabel("log_liklihood")
plt.show()

plt.plot(log_liklihood_v[1:-1])
plt.title("x vs validation ~ log_liklihood_v")
plt.xlabel("Iterations")
plt.ylabel("log_liklihood")
plt.show()


# plt.plot(Rnrm_t[2:-1])
# plt.title("T residuals")
# plt.xlabel("Iterations")
# plt.ylabel("Normalised residuals")
# plt.show()


plt.plot(log_liklihood_T[1:-1])
plt.title("x vs validation ~ log_liklihood_T")
plt.xlabel("Iterations")
plt.ylabel("log_liklihood")
plt.show()

plt.plot(log_liklihood_V[1:-1])
plt.title("x vs validation ~ log_liklihood_V")
plt.xlabel("Iterations")
plt.ylabel("log_liklihood")
plt.show()

plt.plot(psnr_b[1:-1])
plt.title("x vs validation ~ psnr_b")
plt.xlabel("Iterations")
plt.ylabel("psnr")
plt.show()

plt.plot(psnr_V[1:-1])
plt.title("x vs validation ~ psnr_V")
plt.xlabel("Iterations")
plt.ylabel("psnr")
plt.show()

plt.plot(psnr_T[1:-1])
plt.title("x vs validation ~ psnr_T")
plt.xlabel("Iterations")
plt.ylabel("psnr")
plt.show()

# # %%
# # compute second derivative
# from scipy.ndimage import gaussian_filter1d

# smooth = gaussian_filter1d(Rnrm_v, 10)
# smooth = Rnrm_v[20:-1]
# plt.plot(smooth)
# plt.show()
# #  %%
# diff = np.diff(np.diff(smooth,axis=0),axis=0)[1:-2];
# plt.plot(diff)
# plt.show()
# #  %%
# inf_point = np.where ( Rnrm_v == Rnrm_v.max())[0][0]
# inf_point
# plt.plot(smooth)
# plt.vlines(x=inf_point,ymin=smooth.min(),ymax=smooth.max())
# plt.show()
# # find switching points
# # infls = np.where(np.diff(np.sign(diff)))[0]

# # %%

# %%