viewer.py

import numpy as np
import matplotlib.pyplot as plt
from astrodendro import Dendrogram
from astropy.io import fits

class SimpleDendrogramViewer(object):
    
    def __init__(self, array, dendrogram):

        if array.ndim != 2:
            raise ValueError("Only 2-dimensional arrays are supported at this time")

        # Get the lines as individual elements, and the mapping from line to structure
        self.lines, self.mapping = d.get_lines(collection=False, reverse=True)

        # Define inverse mapping from structures to lines
        self.inverse_mapping = {}
        for line in self.mapping:
            s = self.mapping[line]
            if s not in self.inverse_mapping:
                self.inverse_mapping[s] = [line]
            else:
                self.inverse_mapping[s] += [line]

        # Define the list of currently selected structures
        self.selected = []

        # Initiate plot
        self.fig = plt.figure(figsize=(14, 8))
        self.ax1 = self.fig.add_axes([0.1, 0.1, 0.4, 0.8])
        self.ax1.imshow(image, origin='lower')
        self.ax2 = self.fig.add_axes([0.6, 0.3, 0.35, 0.4])
        xmin, xmax = np.inf, -np.inf
        ymin, ymax = np.inf, -np.inf
        for line in self.lines:
            x, y = line.get_xdata(), line.get_ydata()
            xmin = min(xmin, np.min(x))
            xmax = max(xmax, np.max(x))
            ymin = min(ymin, np.min(y))
            ymax = max(ymax, np.max(y))
            line.set_picker(True)
            self.ax2.add_line(line)
        dx = xmax - xmin
        self.ax2.set_xlim(xmin - dx * 0.1, xmax + dx * 0.1)
        self.ax2.set_ylim(ymin * 0.5, ymax * 2.0)
        self.ax2.set_yscale('log')

        self.prev_contour = None

        self.fig.canvas.mpl_connect('pick_event', self.line_picker)
        
        plt.show()

    def line_picker(self, event):
        if self.prev_contour is not None:
            for collection in self.prev_contour.collections:
                self.ax1.collections.remove(collection)
        for l in self.selected:
            l.set_color('blue')
            l.set_lw(1)
        self.selected = []
        structure = self.mapping[event.artist]
        for s in structure.descendants + [structure]:
            for l in self.inverse_mapping[s]:
                l.set_color('red')
                l.set_lw(2)
                self.selected.append(l)
        self.prev_contour = self.ax1.contour(structure.get_mask(image.shape, subtree=True), colors='red', linewidths=2, levels=[0.5])
        event.canvas.draw()

image = fits.getdata("MSX_E.fits")
d = Dendrogram.compute(image, verbose=True, min_data_value=3.e-5)

v = SimpleDendrogramViewer(image, d)