Add time lag analysis helpers (#109)

* add time lag analysis helpers

* revert channels as static property

Author: wtbarnes

pyproject.toml

@@ -30,8 +30,9 @@ Homepage = "https://github.com/wtbarnes/synthesizAR"
 Documentation = "https://synthesizar.readthedocs.io"
 
 [project.optional-dependencies]
-all = ["synthesizAR[aia,atomic,ebtel,hydrad,parallel,xrt]"]
+all = ["synthesizAR[aia,analysis,atomic,ebtel,hydrad,parallel,xrt]"]
 aia = ["aiapy"]
+analysis = ["sunkit_image"]
 atomic = [
   "plasmapy",
   "fiasco",

synthesizAR/analysis/time_lag/__init__.py

@@ -0,0 +1,4 @@
+"""
+Convenience functions for computing and analyzing time lag measurements from datacubes.
+"""
+from .time_lag import *

synthesizAR/analysis/time_lag/map_sources.py

@@ -0,0 +1,49 @@
+"""
+Map sources for diagnostic maps.
+"""
+from astropy.visualization import ImageNormalize
+from sunpy.map import GenericMap
+
+import synthesizAR.visualize.colormaps  # NOQA
+
+
+class TimeLagMap(GenericMap):
+    """
+    A map that represents the time lag between two images.
+    """
+    def __init__(self, data, header, **kwargs):
+        super().__init__(data, header, **kwargs)
+        self.plot_settings['cmap'] = 'idl_bgry_004'
+        self.plot_settings['norm'] = ImageNormalize(vmin=-7200, vmax=7200)
+        self.nickname = f"{header.get('chan_a')}-{header.get('chan_b')}"
+
+    @property
+    def measurement(self):
+        return ' '.join(self.meta.get('measrmnt', '').split('_')).capitalize()
+
+    @classmethod
+    def is_datasource_for(cls, data, header, **kwargs):
+        return ('chan_a' in header and
+                'chan_b' in header and
+                header.get('measrmnt')=='time_lag')
+
+
+class CrossCorrelationMap(GenericMap):
+    """
+    A map that represents the maximum cross-correlation value between two images
+    """
+    def __init__(self, data, header, **kwargs):
+        super().__init__(data, header, **kwargs)
+        self.plot_settings['cmap'] = 'magma'
+        self.plot_settings['norm'] = ImageNormalize(vmin=0, vmax=1)
+        self.nickname = f"{header.get('chan_a')}-{header.get('chan_b')}"
+
+    @property
+    def measurement(self):
+        return ' '.join(self.meta.get('measrmnt', '').split('_')).capitalize()
+
+    @classmethod
+    def is_datasource_for(cls, data, header, **kwargs):
+        return ('chan_a' in header and
+                'chan_b' in header and
+                header.get('measrmnt')=='max_cross_correlation')

synthesizAR/analysis/time_lag/time_lag.py

@@ -0,0 +1,85 @@
+"""
+Wrapper functions for producing time lag and cross-correlation maps from data cubes.
+"""
+import astropy.units as u
+import itertools
+import sunkit_image.time_lag
+import sunpy.map
+
+from synthesizAR.instruments.sdo import _AIA_CHANNEL_WAVELENGTHS
+
+import synthesizAR.analysis.time_lag.map_sources  # NOQA
+
+
+__all__ = ['get_aia_channel_combinations', 'make_time_lag_map', 'make_cross_correlation_map']
+
+
+def get_aia_channel_combinations():
+    """
+    Convenience function for listing all possible AIA channel pairs.
+    This is useful for computing time lags.
+    """
+    channel_list = [f"{chan.to_value('Angstrom'):.0f}" for chan in _AIA_CHANNEL_WAVELENGTHS]
+    channel_combinations = list(itertools.combinations(channel_list, 2))
+    channel_combinations = channel_combinations[:5] + [sorted(c, key=lambda x: float(x), reverse=True) for c in channel_combinations[5:]]
+    return channel_combinations
+
+
+def _get_meta_and_time(cube_a, cube_b, lag_bounds):
+    time_a = cube_a.axis_world_coords('time')[0]
+    time_b = cube_b.axis_world_coords('time')[0]
+    if not (time_a == time_b).all():
+        raise ValueError('Time axes of both cubes must be the same')
+    time = (time_a - time_a[0]).to('s')
+    if lag_bounds is None:
+        lag_bounds = u.Quantity([-time[-1]/2, time[-1]/2])
+    meta = cube_a.meta.copy()
+    stale_keys = ['bunit', 'date_sim', 'wavelnth', 'waveunit', 'instrume', 'telescop', 'obsrvtry', 'detector']
+    for k in stale_keys:
+        _ = meta.pop(k)
+    meta['chan_a'] = cube_a.meta.get('wavelnth')
+    meta['chan_b'] = cube_b.meta.get('wavelnth')
+    return time, lag_bounds, meta
+
+
+@u.quantity_input
+def make_time_lag_map(cube_a, cube_b, lag_bounds: u.s=None):
+    """
+    Coordinate-aware wrapper around `~sunkit_image.time_lag.time_lag`
+
+    Parameters
+    ----------
+    cube_a : `~ndcube.NDCube`
+    cube_b : `~ndcube.NDCube`
+    lag_bounds : `~astropy.units.Quantity`
+
+    Return
+    ------
+    : `~sunpy.map.GenericMap`
+    """
+    time, lag_bounds, meta = _get_meta_and_time(cube_a, cube_b, lag_bounds)
+    data = sunkit_image.time_lag.time_lag(cube_a.data, cube_b.data, time, lag_bounds=lag_bounds)
+    meta['bunit'] = time.unit.to_string(format='FITS')
+    meta['measrmnt'] = 'time_lag'
+    return sunpy.map.Map(data, meta)
+
+
+@u.quantity_input
+def make_cross_correlation_map(cube_a, cube_b, lag_bounds: u.s=None):
+    """
+    Coordinate-aware wrapper around `~sunkit_image.time_lag.max_cross_correlation`
+
+    Parameters
+    ----------
+    cube_a : `~ndcube.NDCube`
+    cube_b : `~ndcube.NDCube`
+    lag_bounds : `~astropy.units.Quantity`
+
+    Return
+    ------
+    : `~sunpy.map.GenericMap`
+    """
+    time, lag_bounds, meta = _get_meta_and_time(cube_a, cube_b, lag_bounds)
+    data = sunkit_image.time_lag.max_cross_correlation(cube_a.data, cube_b.data, time, lag_bounds=lag_bounds)
+    meta['measrmnt'] = 'max_cross_correlation'
+    return sunpy.map.Map(data, meta)

synthesizAR/instruments/sdo.py

@@ -4,7 +4,10 @@
 """
 import asdf
 import astropy.units as u
+import ndcube
 import numpy as np
+import pathlib
+import sunpy.map
 
 from aiapy.psf import filter_mesh_parameters
 from aiapy.response import Channel
@@ -13,6 +16,7 @@
 from scipy.interpolate import interpn
 
 from synthesizAR.instruments import InstrumentBase
+from synthesizAR.instruments.util import map_list_to_time_cube
 from synthesizAR.util.decorators import return_quantity_as_tuple
 
 __all__ = ['InstrumentSDOAIA', 'aia_kernel_quick']
@@ -22,6 +26,15 @@
 with asdf.open(_TEMPERATURE_RESPONSE_FILE, 'r', memmap=False) as af:
     _TEMPERATURE_RESPONSE = af.tree
 
+_AIA_CHANNEL_WAVELENGTHS = [
+    94*u.angstrom,
+    131*u.angstrom,
+    171*u.angstrom,
+    193*u.angstrom,
+    211*u.angstrom,
+    335*u.angstrom,
+]
+
 
 class AIAChannel(Channel):
 
@@ -72,6 +85,10 @@ def __init__(self, observing_time, observer, **kwargs):
             **kwargs,
         )
 
+    @cached_property
+    def channels(self):
+        return [AIAChannel(w) for w in _AIA_CHANNEL_WAVELENGTHS]
+
     @property
     def observatory(self):
         return 'SDO'
@@ -84,17 +101,6 @@ def detector(self):
     def telescope(self):
         return 'SDO/AIA'
 
-    @cached_property
-    def channels(self):
-        return [
-            AIAChannel(94*u.angstrom),
-            AIAChannel(131*u.angstrom),
-            AIAChannel(171*u.angstrom),
-            AIAChannel(193*u.angstrom),
-            AIAChannel(211*u.angstrom),
-            AIAChannel(335*u.angstrom),
-        ]
-
     @property
     def _expected_unit(self):
         return u.DN / (u.pix * u.s)
@@ -136,6 +142,30 @@ def calculate_intensity_kernel(loop, channel, **kwargs):
             kernel = aia_kernel_quick(channel.name, loop.electron_temperature, loop.density)
         return kernel
 
+    @staticmethod
+    def build_collection_from_maps(save_directory, channels=None):
+        """
+        Build an `~ndcube.NDCollection` from all maps in all channels produced by a simulation.
+
+        Parameters
+        ----------
+        save_directory : path-like
+            Directory that contains all FITS files produced by the `~synthesizAR.instruments.InstrumentBase.observe` method.
+        channels : `list` of `str`
+            Channel labels from which to build collection. If not specified, will default to the
+            six EUV channels of AIA.
+        """
+        save_directory = pathlib.Path(save_directory)
+        if channels is None:
+            channels = [f"{chan.to_value('Angstrom'):.0f}" for chan in _AIA_CHANNEL_WAVELENGTHS]
+        key_data_pairs = []
+        for chan in channels:
+            filenames = sorted(save_directory.glob(f'm_{chan}_t*.fits'))
+            cube = map_list_to_time_cube(sunpy.map.Map(filenames))
+            key_data_pairs.append((chan, cube))
+        col = ndcube.NDCollection(key_data_pairs, aligned_axes=(0, 1, 2))
+        return col
+
 
 @u.quantity_input
 def aia_kernel_quick(channel,

synthesizAR/instruments/util.py

@@ -8,13 +8,14 @@
 import numpy as np
 import xarray
 
-from ndcube.extra_coords.table_coord import QuantityTableCoordinate
+from ndcube.extra_coords.table_coord import QuantityTableCoordinate, TimeTableCoordinate
 from ndcube.wcs.wrappers import CompoundLowLevelWCS
 
 __all__ = [
     'get_wave_keys',
     'add_wave_keys_to_header',
     'extend_celestial_wcs',
+    'map_list_to_time_cube',
     'read_cube_from_dataset',
     'write_cube_to_netcdf',
 ]
@@ -58,11 +59,28 @@ def extend_celestial_wcs(celestial_wcs, *extra_coords, **kwargs):
     for ec in extra_coords:
         if isinstance(ec, tuple):
             array, name, physical_type = ec
-            ec = QuantityTableCoordinate(array, names=name, physical_types=physical_type)
+            if isinstance(array, u.Quantity):
+                ec = QuantityTableCoordinate(array, names=name, physical_types=physical_type)
+            elif isinstance(array, astropy.time.Time):
+                ec = TimeTableCoordinate(array, names=name, physical_types=physical_type)
+            else:
+                raise TypeError(f'{array} has unrecognized type {type(array)}')
         wcses.append(ec.wcs)
     return CompoundLowLevelWCS(celestial_wcs, *wcses, **kwargs)
 
 
+def map_list_to_time_cube(map_list):
+    """
+    Transform a list of maps into a single `ndcube.NDCube`
+    """
+    data = np.array([m.data for m in map_list])
+    # NOTE: This is specific to datacubes produced by synthesizAR
+    times = astropy.time.Time([m.meta['DATE_SIM'] for m in map_list])
+    new_wcs = extend_celestial_wcs(map_list[0].wcs, (times, 'time', 'time'))
+    new_meta = map_list[0].meta.copy()
+    return ndcube.NDCube(data, wcs=new_wcs, meta=new_meta, unit=map_list[0].unit)
+
+
 def read_cube_from_dataset(filename, axis_name, physical_type):
     """
     Read an `~ndcube.NDCube` from an `xarray` dataset.