First commit

import pickle

import plotly.graph_objects as go

import matplotlib.pyplot as plt

import numpy as np

source_names = ['G332', 'G28', 'G14', 'G12']

source = source_names[0]

with open(f'final_results_{source}_fullres.pickle', 'rb') as infile:

    fit_results = pickle.load(infile)

rms = 0.003

detection_86 = fit_results[0,:,:] > 3 * rms

detection_A = fit_results[3,:,:] > 3 * rms

discard_failed_86 = (fit_results[0,:,:] < 0.1) & (fit_results[2,:,:] > 0.3)

discard_failed_87 = fit_results[4,:,:] < 0.1

discard_failed_88 = fit_results[5,:,:] < 0.1

discard_failed = discard_failed_86 & discard_failed_87 & discard_failed_88

unsaturated_87_86 = fit_results[4,:,:] / fit_results[0,:,:] < 1.2

velocity_field_1 = np.where(detection_A & discard_failed_86 & (np.abs(fit_results[1,:,:]) < 200), fit_results[1,:,:], np.nan)

velocity_field_2 = np.where(detection_A & discard_failed_86 & (np.abs(fit_results[7,:,:]) < 200), fit_results[7,:,:], np.nan)

velocity_field_3 = np.where(detection_A & discard_failed_86 & (np.abs(fit_results[13,:,:]) < 200), fit_results[13,:,:], np.nan)

ncomps = np.nansum(~np.isnan(fit_results[(0, 6, 12),:,:]), axis=0)

peak_component = np.nanmax(fit_results[(0, 6, 12),:,:], axis=0)

comp_idx = np.argmax(fit_results[(0, 6, 12),:,:] == peak_component, axis=0)

velocity_field = np.where((comp_idx == 0) & (fit_results[1,:,:] > -200), fit_results[1,:,:], velocity_field_1)

velocity_field = np.where((comp_idx == 1) & (fit_results[7,:,:] > -200), fit_results[7,:,:], velocity_field)

velocity_field = np.where((comp_idx == 2) & (fit_results[13,:,:] > -200), fit_results[13,:,:], velocity_field)

velocity_field = np.where(velocity_field > -200, velocity_field, np.nan)

grid = np.meshgrid(np.arange(512),np.arange(512))

fig = go.Figure(data=[go.Scatter3d(x=grid[0].flatten(), y=grid[1].flatten(),z=velocity_field_1.flatten(), mode='markers', opacity=0.1, marker=dict(size=1,symbol='cross'))])

fig.add_scatter3d(x=grid[0].flatten(), y=grid[1].flatten(),z=velocity_field_2.flatten(), mode='markers', opacity=0.1, marker=dict(size=1,symbol='cross'))

fig.add_scatter3d(x=grid[0].flatten(), y=grid[1].flatten(),z=velocity_field_3.flatten(), mode='markers', opacity=0.1, marker=dict(size=1,symbol='cross'))

fig = go.Figure(data=[go.Surface(z=velocity_field)])

ratio_1 = np.where(detection_86 & discard_failed & unsaturated_87_86,

                   fit_results[4,:,:] / fit_results[0,:,:],

                   np.nan)

ratio_2 = np.where((fit_results[6,:,:] > 3 * 0.003) & (fit_results[6,:,:] < 1) & (fit_results[10,:,:] < 1) & (fit_results[10,:,:] / fit_results[6,:,:] < 1.1), fit_results[10,:,:] / fit_results[6,:,:], np.nan)

ratio_3 = np.where((fit_results[12,:,:] > 3 * 0.003) & (fit_results[12,:,:] < 1) & (fit_results[16,:,:] < 1) & (fit_results[16,:,:] / fit_results[12,:,:] < 1.1), fit_results[16,:,:] / fit_results[12,:,:], np.nan)

ratio_map = np.where((comp_idx == 0) & detection_86 & discard_failed & unsaturated_87_86, fit_results[4,:,:] / fit_results[0,:,:], np.nan)

ratio_map = np.where((comp_idx == 1) & (fit_results[6,:,:] > 3 * 0.003) & discard_failed & (fit_results[10,:,:] / fit_results[6,:,:] < 1.2), fit_results[10,:,:] / fit_results[6,:,:], ratio_map)

ratio_map = np.where((comp_idx == 2) & (fit_results[12,:,:] > 3 * 0.003) & discard_failed & (fit_results[16,:,:] / fit_results[12,:,:] < 1.2), fit_results[16,:,:] / fit_results[12,:,:], ratio_map)

ratio_map = np.where(fit_results[0,:,:] < fit_results[6,:,:], ratio_2, ratio_1)

density_map = 10 ** ((np.arctanh(ratio_map - 0.55) + 2.3) ** 0.9 + 3.75)

density_map = np.where(density_map < 5e4, 5e4, density_map)

size_arr = np.where(~np.isnan(density_map), fit_results[0,:,:], 0)

fig = go.Figure(data=[go.Surface(z=density_map, surfacecolor=size_arr)])

# grid = np.meshgrid(np.arange(512),np.arange(512))

# fig = go.Figure(data=[go.Scatter3d(x=grid[0].flatten(), y=grid[1].flatten(),z=density_map.flatten(), mode='markers', opacity=0.3, marker=dict(size=size_arr, symbol='cross'))])

x_eye = -1.25

y_eye = 2

z_eye = 2

fig.update_layout(

         title=None,

         width=600,

         height=600,

         paper_bgcolor='rgba(0,0,0,0)',

         plot_bgcolor='rgba(0,0,0,0)',

         font_color="white",

         scene=dict(zaxis=dict(type='log', dtick=1)),

         margin=dict(l=0, r=0, b=0, t=0),

         modebar = dict(bgcolor='rgba(0, 0, 0, 0)'),

         scene_camera_eye=dict(x=x_eye, y=y_eye, z=z_eye),

         updatemenus=[dict(type='buttons',

                  showactive=False,

                  y=1,

                  x=1,

                  xanchor='left',

                  yanchor='bottom',

                  pad=dict(t=0, r=0),

                  buttons=[dict(label='Play',

                                method='animate',

                                visible=False,

                                args=[None, dict(frame=dict(duration=5, redraw=True),

                                                             transition=dict(duration=0),

                                                             fromcurrent=True,

                                                             mode='immediate'

                                                            )]

                                            )

                                      ]

                              )

                        ]

)

def rotate_z(x, y, z, theta):

    w = x+1j*y

    return np.real(np.exp(1j*theta)*w), np.imag(np.exp(1j*theta)*w), z

frames=[]

for t in np.arange(0, 6.26, 0.1):

    xe, ye, ze = rotate_z(x_eye, y_eye, z_eye, -t)

    frames.append(go.Frame(layout=dict(scene_camera_eye=dict(x=xe, y=ye, z=ze))))

fig.frames=frames

plotly.offline.plot(fig, filename=r'test_animation.html')

fig.show()