get_clust_chart.py

import matplotlib.pyplot as plt
import numpy as np


def plt_f_chart(fig, x_name, y_name, coord, x_zmin, x_zmax, y_zmin,
                y_zmax, x_data, y_data, st_sizes_arr, labels, center_cl,
                clust_rad, mps):
    '''
    Zoom on x,y finding chart.
    '''
    ax = fig.add_subplot(111, aspect='equal')
    # Force square plot.
    ax.set_aspect('equal')
    # Set plot limits.
    plt.xlim(x_zmin, x_zmax)
    plt.ylim(y_zmin, y_zmax)
    if coord == 'deg':
        # If RA is used, invert axis.
        ax.invert_xaxis()
    # Set axis labels
    plt.xlabel('{} ({})'.format(x_name, coord), fontsize=30)
    plt.ylabel('{} ({})'.format(y_name, coord), fontsize=30)
    plt.tick_params(axis='both', which='major', labelsize=25)
    # Set minor ticks
    ax.minorticks_on()
    # Background color.
    ax.patch.set_facecolor('#F3F3F3')
    # Set grid
    ax.grid(b=True, which='major', color='gray', linestyle='--', lw=1)
    # Plot radius.
    circle = plt.Circle((center_cl[0], center_cl[1]), clust_rad, color='r',
                        fill=False, lw=1.5, zorder=5)
    fig.gca().add_artist(circle)
    # Color map, higher prob stars look redder.
    cm = plt.cm.get_cmap('RdYlBu_r')
    # If stars have a range of colors, use list of colors. Else use a single
    # color.
    v_min_mp, v_max_mp = round(min(mps), 2), round(max(mps), 2)
    if v_min_mp != v_max_mp:
        st_col = mps
    else:
        st_col = '#4682b4'
    # Plot stars.
    plt.scatter(x_data, y_data, marker='o', c=st_col, s=st_sizes_arr,
                edgecolors='black', cmap=cm, lw=0.25, vmin=v_min_mp,
                vmax=v_max_mp, zorder=3)
    # Plot IDs
    for label, x, y in zip(labels, x_data, y_data):
        plt.annotate(
            label, xy=(x, y),
            xytext=(4, 4),
            textcoords='offset points',
            ha='left', va='bottom',
            fontsize=20, color='#1A5233', zorder=4)


def random_points():
    mu, sigma = np.random.uniform(100, 1000), np.random.uniform(0., 100)
    return mu, sigma


def frame_zoomed(x_data, y_data, center_cl, clust_rad):
    '''
    If possible, define zoomed frame.
    '''

    x_min, x_max = min(x_data), max(x_data)
    y_min, y_max = min(y_data), max(y_data)

    x_zmin, x_zmax = max(x_min, (center_cl[0] - 1.1 * clust_rad)), \
        min(x_max, (center_cl[0] + 1.1 * clust_rad))
    y_zmin, y_zmax = max(y_min, (center_cl[1] - 1.1 * clust_rad)), \
        min(y_max, (center_cl[1] + 1.1 * clust_rad))

    # Prevent axis stretching.
    if (x_zmax - x_zmin) != (y_zmax - y_zmin):
        lst = [(x_zmax - x_zmin), (y_zmax - y_zmin)]
        val, idx = min((val, idx) for (idx, val) in enumerate(lst))
        if idx == 0:
            x_zmax = x_zmin + lst[1]
        else:
            y_zmax = y_zmin + lst[0]

    return x_zmin, x_zmax, y_zmin, y_zmax


def star_size(mag_data):
    '''
    Convert magnitudes into intensities and define sizes of stars in
    finding chart.
    '''
    # 'brightness' is the b1/b2 ratio were b2 is the minimum magnitude value.
    brightness = 10 ** ((np.array(mag_data) - min(mag_data)) / -2.5)
    factor = 1500.
    return 50. + factor * brightness


def main():

    # some random points
    N = 200
    mux, sigmax = random_points()
    muy, sigmay = random_points()
    x = np.random.normal(mux, sigmax, N)
    y = np.random.normal(muy, sigmay, N)
    mag = np.random.uniform(10, 22, N)
    mag_bright = np.random.uniform(9, 11, int(N * 0.1))
    mag_rest = np.random.uniform(11, 22, int(N * 0.9))
    mag = np.append(mag_bright, mag_rest)
    mps = np.random.uniform(0., 1., N)
    labels = [str(_) for _ in np.arange(0, N)]

    clust_rad = (sigmax + sigmay) / 2.
    center_cl = [mux, muy]

    x_zmin, x_zmax, y_zmin, y_zmax = frame_zoomed(x, y, center_cl, clust_rad)

    # Sizes of stars related to their magnitudes.
    st_sizes_arr = star_size(mag)

    x_name, y_name, coord = 'x', 'y', 'px'

    # Plot.
    fig = plt.figure(figsize=(30, 25))  # create the top-level container

    plt_f_chart(fig, x_name, y_name, coord, x_zmin, x_zmax, y_zmin,
                y_zmax, x, y, st_sizes_arr, labels, center_cl, clust_rad, mps)

    # Generate output file for each data file.
    fig.tight_layout()
    plt.savefig('finding_chart.png', dpi=150)

    # Close to release memory.
    plt.clf()
    plt.close()


if __name__ == "__main__":
    main()