#

# Load in multiple noisy realizations of a given simulated wave, and

# characterize the performance of AWARE on detecting the wave

#

import os

import pickle

import numpy as np

import matplotlib.pyplot as plt

import astropy.units as u

import matplotlib

# AWARE constants

import aware_constants

# AWARE utilities

import aware_utils

from aware_plot import longitudinal_lines

# Simulated wave parameters

import swave_params

#

from aware_constants import solar_circumference_per_degree_in_km

import swave_study as sws

# Output fontsize

matplotlib.rcParams.update({'font.size': 18})

# Simulated data

# TODO - run the same analysis on multiple noisy realizations of the simulated

# TODO - data. Generate average result plots over the multiple realizations

# TODO - either as relative error or absolute error as appropriate.

# TODO - First mode: use multiple noisy realizations of the same model.

# TODO - Second mode: choose a bunch of simulated data parameters, and then

# TODO - randomly select values for them (within reason).

# TODO - The recovered parameters should be reasonably close to the simulated

# TODO - parameters.

#

#

use_error_bar = False

for_paper = True

zorder_min = 10000000

true_velocity_kwargs = {"color": "blue", "linewidth": 3, "linestyle": "-.", "zorder": zorder_min+1}

true_acceleration_kwargs = {"color": "red", "linewidth": 3, "linestyle": "-.", "zorder": zorder_min+1}

# Define the Analysis object

class Analysis:

self.aware_version = None

###############################################################################

#

# How to select good arcs

#

# Reduced chi-squared must be below this limit

rchi2_limit = 1.0 #sws.rchi2_limit

###############################################################################

#

# Simulated observations of a wave

#

example = sws.wave_name

# Number of trials

ntrials = sws.n_random

# Number of images

max_steps = sws.max_steps

# Use the second version of the HG to HPC transform

use_transform2 = sws.use_transform2

###############################################################################

#

# Preparation of the simulated observations to create a mapcube that will be

# used by AWARE.

#

# Analysis source data

analysis_data_sources = sws.analysis_data_sources

# Summing of the simulated observations in the time direction

temporal_summing = sws.temporal_summing

# Summing of the simulated observations in the spatial directions

spatial_summing = sws.spatial_summing

# HPC to HG transformation: methods used to calculate the griddata interpolation

griddata_methods = sws.griddata_methods

###############################################################################

#

# AWARE processing: details

#

# Which version of AWARE to use?

aware_version = sws.aware_version

# Radii of the morphological operations in the HG co-ordinate and HPC

# co-ordinates

radii = sws.morphology_radii(aware_version)

################################################################################

#

# Measure the velocity and acceleration of the HG arcs

#

# Position measuring choices

position_choice = sws.position_choice

error_choice = sws.error_choice

# Number of degrees in the polynomial fit

n_degrees = sws.n_degrees

# RANSAC

ransac_kwargs = sws.ransac_kwargs

# Great circle points

great_circle_points = sws.great_circle_points

################################################################################

#

# Where to dump the output

#

# Output directory

output = sws.output

# Special designation: an extra description added to the file and directory

# names in order to differentiate between experiments on the same example wave.

# special_designation = '_ignore_first_six_points'

# special_designation = '_after_editing_for_dsun_and_consolidation'

# special_designation = '_fix_for_crpix12'

special_designation = sws.special_designation

# Output types

otypes = sws.otypes

# Save images in this format

image_file_type = sws.image_file_type

###############################################################################

###############################################################################

#

# Everything below here is set from above

#

# Output directories and filename

odir = os.path.expanduser(output)

otypes_dir = {}

otypes_filename = {}

# Morphological radii

sradii = ''

for r in radii:

for v in r:

sradii = sradii + str(v.value) + '_'

sradii = sradii[0: -1]

# Create the storage directories and filenames

for ot in otypes:

# root directory

idir = os.path.join(odir, ot)

# filename

filename = ''

# All the subdirectories

for loc in [example + special_designation,

'use_transform2=' + str(use_transform2),

'finalmaps',

str(ntrials) + '_' + str(max_steps) + '_' + str(temporal_summing) + '_' + str(spatial_summing.value),

sradii,

position_choice + '_' + error_choice,

aware_utils.convert_dict_to_single_string(ransac_kwargs)]:

idir = os.path.join(idir, loc)

filename = filename + loc + '.'

filename = filename[0: -1]

if not(os.path.exists(idir)):

os.makedirs(idir)

otypes_dir[ot] = idir

otypes_filename[ot] = filename + '.' + str(great_circle_points)

# Load in the wave params

if not sws.observational:

params = swave_params.waves()[example]

#

# Load the results

#

if not os.path.exists(otypes_dir['dat']):

os.makedirs(otypes_dir['dat'])

filepath = os.path.join(otypes_dir['dat'], otypes_filename['dat'] + '.pkl')

print('\nLoading ' + filepath + '\n')

f = open(filepath, 'rb')

results = pickle.load(f)

f.close()

# How many arcs?

nlon = len(results[0])

angles = ((np.linspace(0, 2*np.pi, nlon+1))[0:-1] * u.rad).to(u.deg)

# Long score

long_score = np.asarray([aaa[1].answer.long_score.final_score if aaa[1].answer.fitted else 0.0 for aaa in results[0]])

# Best Long score

long_score_argmax = long_score.argmax()

# Initial value to the velocity

velocity_unit = u.km/u.s

acceleration_unit = u.km/u.s/u.s

if not sws.observational:

true_values = {"velocity": (params['speed'][0] * aware_constants.solar_circumference_per_degree).to(velocity_unit).value,

"acceleration": (params['acceleration'] * aware_constants.solar_circumference_per_degree).to(acceleration_unit).value}

true_value_labels = {"velocity": velocity_unit.to_string('latex_inline'),

"acceleration": acceleration_unit.to_string('latex_inline')}

longitudinal_lines_kwargs = {"bbox": dict(facecolor='yellow', alpha=0.8),

}

best_long_score_text_kwargs = {"bbox": dict(facecolor='red', alpha=0.8),

}

# Legend keywords

legend_kwargs = {"framealpha": 0.7, "facecolor": "yellow", "loc": "best", "fontsize": 12}

def extract(results, n_degree=1, measurement_type='velocity'):

return fitted, rchi2, measurement, measurement_error

def summarize(fitted, rchi2, measurement, rchi2_limit=1.5):

("n_found", n_found)

for n_degree in [1, 2]:

if n_degree == 1:

measurement_types = ['velocity']

fit = 'linear fit'

if n_degree == 2:

measurement_types = ['velocity', 'acceleration']

fit = 'quadratic fit'

for measurement_type in measurement_types:

if measurement_type == 'velocity':

figure_label = '(a)'

if measurement_type == 'acceleration':

figure_label = '(b)'

if not sws.observational:

true_value = true_values[measurement_type]

true_value_label = true_value_labels[measurement_type]

# Make plots of the central tendency of the velocity

fitted, rchi2, measurement, measurement_error = extract(results,

n_degree=n_degree,

measurement_type=measurement_type)

summaries = summarize(fitted, rchi2, measurement, rchi2_limit=rchi2_limit)

print(summaries[0][3], summaries[0][4])

for summary in summaries[0:1]:

plt.close('all')

fig, ax = plt.subplots()

# ax.errorbar(angles.value, summary[1], summary[2], linewidth=0.5, color='green', label=summary[0])

ax.plot(angles.value, summary[1], color='green', label='mean')

ax.plot(angles.value, summary[1] + summary[2], color='green', linestyle=":", label='mean $\pm$ standard deviation')

ax.plot(angles.value, summary[1] - summary[2], color='green', linestyle=":")

ax.xaxis.set_ticks(np.arange(0, 360, 45))

ax.grid('on', linestyle=":")

if not sws.observational:

hline_label = "true {:s} ({:n} {:s})".format(measurement_type, true_value, true_value_label)

if measurement_type == 'velocity':

ax.axhline(true_value, label=hline_label, **true_velocity_kwargs)

if measurement_type == 'acceleration':

ax.axhline(true_value, label=hline_label, **true_acceleration_kwargs)

for key in longitudinal_lines.keys():

ax.axvline(key, **longitudinal_lines[key]['kwargs'])

ax.axvline(long_score_argmax, color='red',

label='best Long score (' + str(long_score_argmax) + u.degree.to_string('latex_inline') + ')')

ax.set_xlabel('longitude (degrees)')

ax.set_ylabel(measurement_type + " ({:s})".format(true_value_label))

if for_paper:

title = "{:s} {:s}\n({:s})".format(figure_label, measurement_type, summary[0])

title = "{:s} {:s}".format(figure_label, measurement_type)

else:

title = "{:s} ({:s})\n{:s}".format(measurement_type, summary[0], fit)

ax.set_title(title)

ax.legend(**legend_kwargs)

filename = "{:s}.{:s}.{:s}.{:s}".format(otypes_filename["img"], measurement_type, summary[0], fit)

filename = "{:s}".format(aware_utils.clean_for_overleaf(filename))

filename = "{:s}.{:s}".format(filename, image_file_type)

file_path = os.path.join(otypes_dir['img'], filename)

print('Saving {:s}'.format(file_path))

fig.tight_layout()

fig.savefig(file_path)

#

# Plot and save the best long score arc

##

plt.close('all')

bls_string = (angles[long_score_argmax].to(u.deg))._repr_latex_()

results[0][long_score_argmax][1].answer.plot(title='wave propagation at the best Long score\n(longitude={:s})'.format(bls_string))

plt.tight_layout()

directory = otypes_dir['img']

filename = "{:s}_{:s}".format(otypes_filename['img'], 'arc_with_highest_score')

filename = "{:s}".format(aware_utils.clean_for_overleaf(filename))

filename = "{:s}.{:s}".format(filename, image_file_type)

full_file_path = os.path.join(directory, filename)

plt.savefig(full_file_path)