camera_plot.py

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib import cm
from PIL import Image
import pandas as pd
from scipy.interpolate import griddata, interp1d
import os

def load_camera():
    """
    Loads raytraced camera data from dat file.
    """
    data = np.loadtxt("Output/Camera.dat")
    newt_data = np.loadtxt("Output/Camera_newton.dat")
    return data, newt_data  


def camera_interpolator(value, x, y, values):
    """
    Interpolates camera data.
    """
    xy = np.concatenate((np.array(x)[:, None], np.array(y)[:, None]), axis=1)
    interp = griddata((x, y), value, values,
                      method='linear', fill_value=np.nan)
    return interp

def plot_camera(data, newt_data):
    """
    Plots the camera data.
    """
    mpl.rcParams.update({'text.color' : "black"})
    data = pd.DataFrame(data, columns=['flux', 'g_do', 'g_sd', 
                                       'alpha', 'beta', 'rn', 'phin'])
    data = data.sort_values(by=['y', 'x'],ignore_index=True)
    newt_data = pd.DataFrame(newt_data, columns=['flux', 'g_do', 'g_sd',
                                                 'alpha', 'beta', 'rnn', 'phin'])
    newt_data = newt_data.sort_values(by=['y', 'x'],ignore_index=True)

    degrees = 30
    mueff = np.cos(np.radians(degrees))

    flux_camera = data[["alpha", "beta", "flux"]]
    g_do_camera = data[["alpha", "beta", "g_do"]]
    g_sd_camera = data[["alpha", "beta", "g_sd"]]

    plt.figure(figsize=(10, 6))
    plt.scatter(flux_camera["alpha"], flux_camera["beta"], c=flux_camera['flux'], cmap='hot', s=1)
    plt.xlabel('Alpha Position')
    plt.ylabel('Beta Position')
    plt.colorbar(label='Flux')
    plt.xlim(np.min(flux_camera["alpha"]),np.max(flux_camera["alpha"]))
    plt.ylim(np.min(flux_camera["beta"]),np.max(flux_camera["beta"]))
    plt.gca().set_facecolor('black')  # Set plot area background to black
    plt.title('Flux')
    plt.savefig("Output/Camera_flux_alpha_beta.png")  # Save with black background
    plt.show()
    plt.close()

    plt.figure(figsize=(10, 6))
    plt.scatter(g_do_camera["alpha"], g_do_camera["beta"], c=g_do_camera['g_do'], cmap='hot', s=1)
    plt.xlabel('Alpha Position')
    plt.ylabel('Beta Position')
    plt.colorbar(label='G-factor shift Disk to observer')
    plt.xlim(np.min(g_do_camera["alpha"]),np.max(g_do_camera["alpha"]))
    plt.ylim(np.min(g_do_camera["beta"]),np.max(g_do_camera["beta"]))
    plt.gca().set_facecolor('black')  # Set plot area background to black
    plt.title('G-factor shift Disk to observer')
    plt.savefig("Output/Camera_g_do_alpha_beta.png")  # Save with black background
    plt.show()
    plt.close()

    plt.figure(figsize=(10, 6))
    plt.scatter(g_sd_camera["alpha"], g_sd_camera["beta"], c=g_sd_camera['g_sd'], cmap='hot', s=1)
    plt.xlabel('Alpha Position')
    plt.ylabel('Beta Position')
    plt.colorbar(label='G-factor shift source to disk')
    plt.xlim(np.min(g_sd_camera["alpha"]),np.max(g_sd_camera["alpha"]))
    plt.ylim(np.min(g_sd_camera["beta"]),np.max(g_sd_camera["beta"]))
    plt.gca().set_facecolor('black')  # Set plot area background to black
    plt.title('G-factor shift source to disk')
    plt.savefig("Output/Camera_g_sd_alpha_beta.png")  # Save with black background
    plt.show()
    plt.close()

    norm = mcolors.Normalize(vmin=np.min([data["flux"].min(), newt_data["flux"].min()]), 
                             vmax=np.max([data["flux"].max(), newt_data["flux"].max()]))

    r = np.geomspace(data["rn"].min(), data["rn"].max(), 400)
    theta = np.linspace(data["phin"].min(), data["phin"].max(), 400)
    R, Theta = np.meshgrid(r, theta)

    newt_r = np.geomspace(newt_data["rnn"].min(), newt_data["rnn"].max(), 400)
    newt_theta = np.linspace(newt_data["phin"].min(), newt_data["phin"].max(), 400)
    newt_R, newt_Theta = np.meshgrid(newt_r, newt_theta)

    # Interpolate flux values onto the GR grid
    flux_grid = griddata((data["phin"], data["rn"]), data["flux"], (Theta, R), method='nearest')
    gdo_grid = griddata((data["phin"], data["rn"]), data["g_do"], (Theta, R), method='nearest')
    # Interpolate flux values onto the Newtonian grid
    newt_flux_grid = griddata((newt_data["phin"], newt_data["rnn"]), newt_data["flux"], (newt_Theta, newt_R), method='nearest')
    newt_gdo_grid = griddata((newt_data["phin"], newt_data["rnn"]), newt_data["g_do"], (newt_Theta, newt_R), method='nearest')

    mpl.rcParams.update({'text.color' : "white"})
    fig, ax = plt.subplots(subplot_kw={'projection': 'polar'},figsize=(10,8),tight_layout=True)
    fig.set_facecolor("black")
    c = ax.pcolormesh(Theta, R, flux_grid, cmap='hot', norm=norm)
    ax.set_title("Flux within r=1000Rg subtended on sky")
    ax.grid(False)
    ax.pcolormesh(newt_Theta, newt_R, newt_flux_grid, cmap='hot', norm=norm)
    cbar = fig.colorbar(c, ax=ax, label='Flux')
    cbar.ax.tick_params(colors='white')  # Set tick labels to white
    cbar.set_label('Flux', color='white')  # Set the label text color to white
    ax.set_facecolor('black')
    ax.set_yticklabels([])  # Hide radial tick labels
    ax.set_xticklabels([])  # Hide angular tick labels
    ax.set_aspect(mueff)
    scale_length = 100  # Length of the scale bar in radial units
    scale_theta = 0  # Angle at which to place the scale bar (90 degrees)
    ax.plot([scale_theta, scale_theta], [900, 900+scale_length], color='white', lw=2)  # Vertical scale bar
    ax.text(scale_theta+0.1, scale_length/2 + 900, f'{scale_length} Rg', color='white', ha='center')
    plt.savefig("Output/Camera_flux_polar.png")  # Save with black background
    plt.show()
    plt.close()

    spectrum_data = np.loadtxt("Output/Camera_spectrum.dat")
    spectrum = spectrum_data[:,1]
    model_egrid = spectrum_data[:,0]
    radii = spectrum_data[:,2]
    spectrum = spectrum[radii == 1]
    model_egrid = model_egrid[radii == 1]
    egrid = np.linspace(0.5,10,96)
    new_egrids = np.ones(shape=(len(egrid), 400, 400))
    for i, matrix in enumerate(new_egrids):
        new_egrids[i] = egrid[i]/gdo_grid
    newt_egrids = np.ones(shape=(len(egrid), 400, 400))
    for i, matrix in enumerate(newt_egrids):
        newt_egrids[i] = egrid[i]/newt_gdo_grid
    f = interp1d(model_egrid, spectrum, kind='linear', fill_value="extrapolate")
    spec = f(new_egrids) 
    newt_spec = f(newt_egrids)

    minimum = 0
    maximum = 0
    for frame in range(len(egrid)):
        new_flux = np.array(spec[frame] * flux_grid)
        newt_new_flux = np.array(newt_spec[frame] * newt_flux_grid)
        if np.max(new_flux) > maximum:
            maximum = np.max(new_flux)
        if np.min(new_flux) < minimum or minimum == 0:
            minimum = np.min(new_flux)
        if np.max(newt_new_flux) > maximum:
            maximum = np.max(newt_new_flux)
        if np.min(newt_new_flux) < minimum or minimum == 0:
            minimum = np.min(newt_new_flux)

    if np.isnan(minimum):
        minimum = 1e-7
    print("Min and Max flux values for color scale:", minimum, maximum)
    norm = mcolors.Normalize(vmin=minimum, vmax=maximum)
    enorm = mcolors.LogNorm(vmin=np.min([new_egrids,newt_egrids]), 
                            vmax=np.max([new_egrids,newt_egrids]))
    output_dir = "Output/frames/"
    os.makedirs(output_dir, exist_ok=True)

    data_cube = pd.DataFrame(columns=["theta", "r", "alpha", "beta", "flux", "energy_bin"])

    for frame in range(len(egrid)):
        fig, ax = plt.subplots(subplot_kw={'projection': 'polar'},figsize=(10,8), tight_layout=True)
        fig.set_facecolor("black")
        new_flux = np.array(spec[frame] * flux_grid)
        newt_new_flux = np.array(newt_spec[frame] * newt_flux_grid)
        norm = mcolors.Normalize(vmin=np.min([new_flux,newt_new_flux]), 
                            vmax=np.max([new_flux,newt_new_flux]))
        mesh = ax.pcolormesh(Theta, R, new_flux, cmap='hot', norm=norm)
        ax.pcolormesh(newt_Theta, newt_R, newt_new_flux, cmap='hot', norm=norm)

        frame_data = pd.DataFrame({
        "theta": Theta.ravel(),  # Flatten the Theta grid
        "r": R.ravel(),          # Flatten the R grid
        "flux": new_flux.ravel(),  # Flatten the flux grid
        "alpha": (R*np.sin(Theta)).ravel(),     # Convert R to alpha
        "beta": (-R*np.cos(Theta)*mueff).ravel(),      # Convert R to beta
        "energy_bin": egrid[frame]  # Add the current energy bin
        })
        data_cube = pd.concat([data_cube, frame_data], ignore_index=True)

        frame_data = pd.DataFrame({
        "theta": newt_Theta.ravel(),  # Flatten the Theta grid
        "r": newt_R.ravel(),          # Flatten the R grid
        "flux": newt_new_flux.ravel(),  # Flatten the flux grid
        "alpha": (newt_R*np.sin(newt_Theta)).ravel(),     # Convert R to alpha
        "beta": (-newt_R*np.cos(newt_Theta)*mueff).ravel(),      # Convert R to beta
        "energy_bin": egrid[frame]  # Add the current energy bin
        })
        data_cube = pd.concat([data_cube, frame_data], ignore_index=True)

        ax.set_title("Flux at Energy {:.2f} keV".format(egrid[0]))
        ax.grid(False)
        cbar = fig.colorbar(mesh, ax=ax, label='Flux',)
        cbar.ax.tick_params(colors='white')  # Set tick labels to white
        cbar.set_label('Flux', color='white')  # Set the label text color to white
        ax.set_facecolor('black')
        ax.set_yticklabels([])  # Hide radial tick labels
        ax.set_xticklabels([])  # Hide angular tick labels
        ax.set_aspect(mueff)
        ax.set_title("Flux at Energy {:.2f} keV".format(egrid[frame]))
        # Add a scale bar manually
        scale_length = 100  # Length of the scale bar in radial units
        scale_theta = 0  # Angle at which to place the scale bar (90 degrees)
        ax.plot([scale_theta, scale_theta], [900, 900+scale_length], color='white', lw=2)  # Vertical scale bar
        ax.text(scale_theta+0.1, scale_length/2 + 900, f'{scale_length} Rg', color='white', ha='center')
        plt.savefig(f"{output_dir}frame_{frame:03d}.png")
        plt.close()

        

    print("Saving flux data to CSV...")
    data_cube.to_csv("Output/flux_data.csv", index=False)
    print("Creating GIFs...")
    frames = []
    for frame in range(len(egrid)):
        img = Image.open(f"{output_dir}frame_{frame:03d}.png")
        frames.append(img)

    # Save as GIF
    frames[0].save(f"Output/XRI_BH_{degrees}.gif", save_all=True, 
                   append_images=frames[1:], duration=500, loop=0)


if __name__ == "__main__":
    data, newt_data = load_camera()
    plot_camera(data, newt_data)