plot_first_order_derivatives()

Complexity

Conditions

25

Size

Total Lines

106

Duplication

Lines	0
Ratio	0 %

Importance

Changes	1
Bugs	0	Features	0

def plot_first_order_derivatives(model, label_names=None, scaled=True,
    show_clipped_region=False, colors=None, zorders=None,
    clip_less_than=None, label_wavelengths=None, latex_label_names=None,
    wavelength_regions=None, show_legend=True, **kwargs):

    if wavelength_regions is None:
        wavelength_regions = [(model.dispersion[0], model.dispersion[-1])]

    if label_names is None:
        label_names = model.vectorizer.label_names

    if latex_label_names is None:
        latex_label_names = {}

    fig, axes = plt.subplots(1, len(wavelength_regions), figsize=(15, 3.5))

    if len(label_names) > 1:
        #cmap = cm.LinearSegmentedColormap.from_list(
        #    "inferno", cmaps._inferno_data, len(label_names))
        cmap = plt.cm.get_cmap("Set1", len(label_names))
    else:
        cmap = lambda x: "k"

    if colors is not None:
        cmap = lambda x: colors[x % len(colors)]

    axes = np.array(axes).flatten()

    scales = []
    for i, label_name in enumerate(label_names):

        # First order derivatives are always indexed first.
        index = 1 + model.vectorizer.label_names.index(label_name)
        y = model.theta[:, index]
        #y = y
        if clip_less_than is not None:
            y[np.abs(y) < clip_less_than] = 0

        scale = np.nanmax(np.abs(y)) if scaled else 1.
        y = y / scale

        c = cmap(i)
        zorder = 1
        if zorders is not None: zorder = zorders[i]
        for ax in axes:
            ax.plot(model.dispersion, y, c=c, zorder=zorder,
                label=latex_label_names.get(label_name, label_name))

            if clip_less_than is not None and show_clipped_region:
                ax.axhspan(-clip_less_than/scale, +clip_less_than/scale,
                    xmin=-1, xmax=+2,
                    facecolor=c, edgecolor=c, zorder=-100, alpha=0.1)


    # Plot any wavelengths.
    if label_wavelengths is not None:
        label_yvalue = 1.0
        for label_name, wavelengths in label_wavelengths.items():
            try:
                color = cmap(label_names.index(label_name))
                label = None
            except (IndexError, ValueError):
                color = 'k'
                #for ax in axes:
                #    ax.plot([model.dispersion[0] - 1], [0], c=color, label=latex_label_names.get(label_name, label_name))

            for ax in axes:
                ax.plot(wavelengths, label_yvalue * np.ones_like(wavelengths),
                    "|", markersize=20, markeredgewidth=2, c=color)
    
    for ax, wavelength_region in zip(axes, wavelength_regions):
        ax.set_xlim(wavelength_region)
        if scaled:
            ax.set_ylim(-1.2, 1.2)
        
        if ax.is_first_col():
            if scaled:
                ax.set_ylabel(r"$\theta/{\max|\theta|}$")
            else:
                ax.set_ylabel(r"$\theta$")
        else:
            ax.set_yticklabels([])
    
        ax.xaxis.set_major_locator(MaxNLocator(4))

    xlabel = r"$\lambda$ $({\rm\AA})$"
    if len(wavelength_regions) == 1:
        ax.set_xlabel(xlabel)

    else:
        _show_xlim_changes(fig)
        
        ax = fig.add_axes([0, 0, 1, 1])
        ax.set_axis_off()
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.text(0.5, 0.05, xlabel, rotation='horizontal',
            horizontalalignment='center', verticalalignment='center')

    fig.tight_layout()

    if show_legend:
        axes[0].legend(loc="upper right", ncol=kwargs.get("legend_ncol", len(label_names) % 7), frameon=False)
    fig.subplots_adjust(wspace=0.01, bottom=0.20)

    return fig