Creating a bunch of plots in a PDF file using astrodendro

gistfile1.py

#!/usr/bin/env python
# coding: utf8
import numpy as np
import sys
import subprocess
import matplotlib as mpl
from matplotlib.backends.backend_pdf import PdfPages
from astrodendro import Dendrogram
from astrocube import DataCube

############# Set up parameters ############

fig_series = "D"
cube = DataCube("base.fits", calc_noise_dev=False)

num_runs = 25
min_flux = 0.5
min_npix = 3
min_delta = 0.2
scale = 0.05 # deviation of noise
seed = 86400
np.random.seed(86400) # consistent random seed
highlights = [ ("purple", (81,74,75)),
               ("green", (24,31,19)),
               ("orange", (38,66,56)),
               ("blue", (46,50,56)),
             ]
fig_export_size = (10,7.5) # in inches

############# Set up Matplotlib ############

def make_figure(title):
    subtitle = "Dendrogram parameters: min. flux {f}, min. npix {n:.0f}, min. delta {d}".format(f=min_flux, n=min_npix, d=min_delta)
    mpl_fig = mpl.figure.Figure()
    mpl_fig.set_size_inches(*fig_export_size)
    mpl_fig.suptitle(title, weight='heavy', size='large')
    mpl_fig.suptitle(subtitle, y=0.95)
    mpl_axes = mpl_fig.add_subplot(111)
    mpl.backends.backend_pdf.FigureCanvasPdf(mpl_fig) # set up a canvas for the figure
    return mpl_fig, mpl_axes

combined_figure, combined_axes = make_figure(title=u"{n} Dendrograms from {obj} {line} with noise sigma={sigma}".format(obj=cube.object_name, line=cube.line_name, n=num_runs, sigma=scale))

pdf_outfile = "fig{0}.pdf".format(fig_series)
pp = PdfPages(pdf_outfile)

############# Create our dendrograms #######

orig_data = cube.data

print("Generating true dendrogram...")
true_d = Dendrogram(orig_data, minimum_flux=min_flux, minimum_npix=min_npix, minimum_delta=min_delta, verbose=False)
# Add the true dendrogram as the first page:
true_fig, true_axes = make_figure("Fig {series} - True dendrogram (noiseless)".format(series=fig_series))
d_plot = true_d.make_plot(true_axes, color=(0,0,0,1))
for h_color, h_point in highlights:
    d_plot.create_highlighter(h_color).highlight(true_d.item_at(h_point))
pp.savefig(true_fig)

for i in range(0,num_runs):
    print("Creating figure {i}/{n}".format(i=i+1,n=num_runs))
    print("  >> Adding noise...")
    data = orig_data + np.random.normal(scale=scale, size=orig_data.shape)
    
    print("  >> Generating dendrogram...")
    d = Dendrogram(data, minimum_flux=min_flux, minimum_npix=min_npix, minimum_delta=min_delta, verbose=False)
    print("  >> Generating plot...")
    # Plot this dendrogram as a new page in the PDF
    new_fig, new_axes = make_figure("Fig {series}-{n}".format(series=fig_series, n=i+1))
    d_plot = d.make_plot(new_axes, color=(0,0,0,1))
    for h_color, h_point in highlights:
        d_plot.create_highlighter(h_color, alpha=1).highlight(d.item_at(h_point))
    pp.savefig(new_fig)
    # Plot this dendrogram on the combined PDF:
    c_plot = d.make_plot(combined_axes, color=(0,0,0,0.3)) # colour is semi-transparent black
    for h_color, h_point in highlights:
        c_plot.create_highlighter(h_color, alpha=0.3, line_width_extra=0).highlight(d.item_at(h_point))

# add the true dendrogram overtop of the combined plot:
true_d.make_plot(combined_axes, "red")

############# Export the PDF ###############
pp.savefig(combined_figure)
pp.infodict()['Title'] = "Dendrograms Figure Series {0}".format(fig_series)
pp.close()
pp = None
# mpl_fig.canvas.draw()