# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:synthetic-observables]
#     language: python
#     name: conda-env-synthetic-observables-py
# ---

# # Compute Timelags

# In this notebook, we'll compute the timelags and cross-correlation for every channel pair and heating type and save them as FITS files.

# +
import os
import sys

import numpy as np
import distributed
import matplotlib.pyplot as plt
import matplotlib.colors
from sunpy.map import Map,GenericMap
from astropy.coordinates import SkyCoord
import astropy.units as u
from astropy.utils.console import ProgressBar

import synthesizAR
from synthesizAR.instruments import InstrumentSDOAIA
from synthesizAR.analysis import DistributedAIACube,AIATimelags
from synthesizAR.visualize import bgry_004_idl_cmap

# %matplotlib inline
# -

synthesizAR.version.version

# Spin up a local Dask cluster.

cluster = distributed.LocalCluster(n_workers=32,threads_per_worker=2,)
client = distributed.Client(cluster)
client

channels = [94,131,171,193,211,335]
channel_pairs = [(94,335),
                 (94,171),
                 (94,193),
                 (94,131),
                 (94,211),
                 (335,131),
                 (335,193),
                 (335,211), 
                 (335,171),
                 (211,131),
                 (211,171),
                 (211,193),
                 (193,171),
                 (193,131),
                 (171,131),]

heating = ['high_frequency', 'intermediate_frequency', 'low_frequency', 'cooling_outofphase_long', 'cooling']
labels = ['High', 'Intermediate', 'Low', 'Random', 'Cooling']

intensity_file_format = '/storage-home/w/wtb2/data/timelag_synthesis_v2/{}/nei/SDO_AIA/{}/map_t{:06d}.fits'

result_file_format = '/storage-home/w/wtb2/projects/synthetic-observables-paper-models/paper/data/{}/{}_{}_{}.fits'

timelag_bounds = (-6*u.hour,6*u.hour)

# Now, compute the timelag and correlation maps for all of the channel pairs and all of the different heating models.

# ## High Frequency

tl = AIATimelags(*[DistributedAIACube.from_files([intensity_file_format.format('high_frequency', c, i) for i in range(500,2500)]) for c in channels])

for ca,cb in channel_pairs:
    timelag_map = tl.make_timelag_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    timelag_map.save(result_file_format.format('high_frequency', 'timelag', ca, cb))
    correlation_map = tl.make_correlation_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    correlation_map.save(result_file_format.format('high_frequency', 'correlation', ca, cb))

# ## Intermediate Frequency

tl = AIATimelags(*[DistributedAIACube.from_files([intensity_file_format.format('intermediate_frequency', c, i) for i in range(500,2500)]) for c in channels])

for ca,cb in channel_pairs:
    timelag_map = tl.make_timelag_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    timelag_map.save(result_file_format.format('intermediate_frequency', 'timelag', ca, cb))
    correlation_map = tl.make_correlation_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    correlation_map.save(result_file_format.format('intermediate_frequency', 'correlation', ca, cb))

# ## Low Frequency

tl = AIATimelags(*[DistributedAIACube.from_files([intensity_file_format.format('low_frequency', c, i) for i in range(500,2500)]) for c in channels])

for ca,cb in channel_pairs:
    timelag_map = tl.make_timelag_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    timelag_map.save(result_file_format.format('low_frequency', 'timelag', ca, cb))
    correlation_map = tl.make_correlation_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    correlation_map.save(result_file_format.format('low_frequency', 'correlation', ca, cb))

# ## Random

tl = AIATimelags(*[DistributedAIACube.from_files([intensity_file_format.format('cooling_outofphase_long', c, i) 
                                                  for i in range(500,2500)]) for c in channels])

for ca,cb in channel_pairs:
    timelag_map = tl.make_timelag_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    timelag_map.save(result_file_format.format('random', 'timelag', ca, cb))
    correlation_map = tl.make_correlation_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    correlation_map.save(result_file_format.format('random', 'correlation', ca, cb))

# ## Cooling

tl = AIATimelags(*[DistributedAIACube.from_files([intensity_file_format.format('cooling', c, i) 
                                                  for i in range(0,1000)]) for c in channels])

for ca,cb in channel_pairs:
    timelag_map = tl.make_timelag_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    timelag_map.save(result_file_format.format('cooling', 'timelag', ca, cb))
    correlation_map = tl.make_correlation_map(f'{ca}',f'{cb}', timelag_bounds=timelag_bounds,chunks=(tl[0].shape[1]//3,tl[0].shape[2]//3))
    correlation_map.save(result_file_format.format('cooling', 'correlation', ca, cb))



StarcoderdataJupyterScriptsDedupFiltered