utils.py

import numpy as np
import numpy.random
import matplotlib
import matplotlib.pyplot as plt

local_rng = numpy.random.default_rng()

# MARK: sampling utilities


def make_grid(limits=(-1, 1), num_divisions=10, count=2):
    """
    Generate a grid of all value combinations for `count` variables
    (parameters or features) where each ranges over the interval
    specified by `limits` and has `num_divisions` evenly-spaced steps.

    NB: the last dimension indexes the variables

    # Arguments
        limits: a tuple (low, high) specifying the endpoints of the range (same for all vars)
        num_divisions: how many steps to divide the range into (including endpoints)
        count: the number of (identically-ranged) variables

    # Returns
        X: a (count + 1)-dimensional array, of size num_divisions on the first count dims
           and size count on the last
    """
    ticks = np.linspace(limits[0], limits[1], num=num_divisions)
    return np.moveaxis(np.stack(np.meshgrid(*((ticks,) * count), indexing="ij")), 0, -1)


def make_random(num_samples=100, limits=(-1, 1), rng=local_rng, count=2):
    """
    Generate a 2D array of uniformly-distributed randomised variable values.

    NB: the last dimension (of 2) indexes the variables

    # Arguments
        num_samples: number of rows to generate
        limits: tuple (low, high) specifying range of values for all vars
        rng: an instance of numpy.random.Generator from which to draw random values
        count: how many vars to generate per sample

    # Returns:
        X: a 2 dimensional array of size `num_samples` x `count`
            (ie, columns are vars and rows are samples)
    """
    scale = limits[1] - limits[0]
    offset = limits[0]
    return rng.random((num_samples, count)) * scale + offset


def grid_sample(
    function, count=2, num_divisions=10, limits=(-1, 1), rng=local_rng, noise=0
):
    """
    Sample a function on a grid, with optional additive Gaussian noise.

    # Arguments
        function: the function to sample. it must be able to be called with a
            single ndarray argument (in which `count` input variables are indexed
            on the last dimension) and return an appropriately sized array of results.
        count: the number of variables the function expects in its input array
        num_divisions: the number of sample points along each variable dimension
        limits: tuple (low, high) specifying range of values for all vars
        rng: an instance of numpy.random.Generator from which to draw random values
        noise: standard deviation of additive Gaussian noise

    # Returns
        X: the grid of variable values (indexed on the last dimension)
        y: an array of the corresponding function values
    """
    X = make_grid(limits, num_divisions, count)
    y = reshaped_apply(X, function)

    if noise > 0:
        y += rng.normal(scale=noise, size=y.shape)

    return X, y


def random_sample(
    function, count=2, num_samples=100, limits=(-1, 1), rng=local_rng, noise=0
):
    """
    Sample a function at random variable values, with optional additive Gaussian noise.

    # Arguments
        function: the function to sample. it must be able to be called with a single
            ndarray argument (in which `count` input variables are indexed on the last
            dimension) and return an appropriately sized array (vector) of results.
        count: the number of variables the function expects in its input array
        num_samples: the total number of points to sample
        limits: tuple (low, high) specifying range of values for all vars
        rng: an instance of numpy.random.Generator from which to draw random values
        noise: standard deviation of additive Gaussian noise

    # Returns
        X: a 2D array of variable values, where rows are samples and columns are variables
        y: an array (vector) of the corresponding function values
    """
    X = make_random(num_samples, limits, rng, count)
    y = function(X)

    if noise > 0:
        y += rng.normal(scale=noise, size=y.shape)

    return X, y


def reshaped_apply(X, func):
    """
    Apply a function to a data array.

    If X has more than 2 dimensions, all but the last are unwound into
    matrix rows, and then the results vector y is rewound back into the
    original shape of those dimensions. The main use case for this is
    to work with grids generated by `make_grid`.

    # Arguments
        X: a table or array of variable values, where the last dimension
           indexes the values
        func: a function taking a single 2d array argument and returning
           a vector

    # Returns
        y: a vector or array of function outputs for the features in X,
           and having a corresponding shape
    """
    rewind = None
    if len(X.shape) > 2:
        rewind = X.shape[:-1]
        X = X.reshape(np.product(rewind), X.shape[-1])

    y = func(X)

    if rewind is not None:
        y = y.reshape(rewind)

    return y


# MARK: linear model utilities


def add_x0(X):
    """
    Prepend a column (or array) of 1s to an array of samples,
    as a dummy feature representing bias/offset/intercept.

    # Arguments
        X: a table or array of variable values, where the last dimension
           indexes the variables

    # Returns:
        X: the table or array with all records prefixed with a constant
           term x0
    """
    x0 = np.expand_dims(np.ones(X.shape[:-1], X.dtype), axis=-1)
    return np.concatenate((x0, X), axis=-1)


def affine(X, weights):
    """
    Generic affine function of X computed as X . [weights]^T

    If the weights vector is 1 longer than the number of features in X, we
    assume it incudes a bias term in first position, and prepend a
    corresponding x0 (=1) to X before multiplying.

    If X has more than 2 dimensions, all but the last are unwound into
    matrix rows, and then the results vector y is rewound back into the
    original shape of those dimensions. The main use case for this is
    to work with grids generated by `make_grid`.

    # Arguments
        X: a table or array of variable values, where the last dimension
           indexes the values
        weights: a vector of coefficients, either the same length as the last
           dimension of X, or 1 longer

    # Returns
        y: a vector or array of dot products of the input observations
           with the weights
    """
    if len(weights) == (X.shape[-1] + 1):
        X = add_x0(X)

    return reshaped_apply(X, lambda z: z @ weights)


# MARK: non-parametric model utilities


def vote(y):
    """
    Find the class with highest number of predictions in
    an array.

    # Arguments
        y: a vector of predictions

    # Returns
        top: the prediction with the highest number of occurrences
            -- ties are won by the numerically lowest class
    """
    classes, counts = np.unique(y, return_counts=True)
    return classes[np.argmax(counts)]


def gini_impurity(y):
    """
    Calculate the Gini impurity for a vector of labels.

    # Arguments
        y: a numpy vector of class labels for all set/node elements

    # Returns
        g: the Gini impurity value for the vector
    """
    return 1 - np.sum([(np.sum(y == cc) / len(y)) ** 2 for cc in np.unique(y)])


# MARK: plotting utilities


def plot_unimplemented(axes, title="Not implemented", msg="Not implemented"):
    """
    Display simple message on plots for unanswered questions.

    # Arguments
        axes: a Matplotlib Axes object into which to plot
        title: title for the plot
        msg: text to display in the plot body

    # Returns
        None
    """
    axes.set_title(title)
    axes.xaxis.set_visible(False)
    axes.yaxis.set_visible(False)
    axes.text(
        0.5,
        0.5,
        msg,
        horizontalalignment="center",
        verticalalignment="center",
        color="grey",
        fontsize="xx-large",
        fontweight="black",
        rotation=30,
    )


def plot_classification_map(
    axes,
    classifier,
    X=None,
    y=None,
    limits=(-5, 5),
    resolution=30,
    title="Classification Map",
    legend_loc="upper left",
):
    """
    Utility function to test a trained classifier on a grid
    and display the "map" of predictions, optionally also
    plotting the training data points if provided.

    # Arguments
        axes: a Matplotlib Axes object into which to plot
        classifier: a classifier function taking a single
            array of sample data in standard X format and
            returning a vector of predicted labels
        X: an array of sample data, where rows are samples
            and the single column is the input feature.
        y: vector of output values corresponding to the
            rows of X
        limits: a tuple (low, high) specifying the value
            range to test for both feature dimensions
        resolution: number of samples across and down the grid
        title: a title for the plot
        legend_loc: where to put the legend, or None to
            omit it

    # Returns
        None
    """

    # colour & marker scheme supports up to 10 classes, though
    # plots will likely be pretty unintelligible with that many
    darks = [plt.cm.tab20.colors[2 * ii] for ii in range(10)]
    lights = [plt.cm.tab20.colors[2 * ii + 1] for ii in range(10)]
    marks = ["o", "v", "*", "x", "+", "s", "D", "^", "h", "2"]

    # this Brewer colour scheme is meant to be accessible, but
    # class 2 light looks very similar to class 1 under some
    # simulated colour vision deficiencies, so we fudge it to be
    # more distinguishable
    lights[2] = (0.7, 1.0, 0.7)

    grid = make_grid(limits=limits, num_divisions=resolution)
    pred_y = reshaped_apply(grid, classifier)

    if y is None:
        nclass = int(np.max(pred_y)) + 1
    else:
        nclass = int(np.max([np.max(pred_y), np.max(y)])) + 1

    fillmap = matplotlib.colors.ListedColormap(lights[:nclass])

    axes.imshow(
        pred_y.T, origin="lower", extent=limits * 2, cmap=fillmap, interpolation="none"
    )

    if (X is not None) and (y is not None):
        for cc in np.unique(y):
            axes.scatter(
                X[y == cc, 0],
                X[y == cc, 1],
                color=darks[cc],
                marker=marks[cc],
                label=cc,
                alpha=0.8,
            )

    axes.set_title(title)
    axes.set_xlabel("$x_1$")
    axes.set_ylabel("$x_2$")

    if legend_loc is not None:
        axes.legend(loc=legend_loc)


def plot_image(axes, image, title=""):
    """
    Convenience wrapper for the imshow settings we pretty much
    *always* want when showing actual images.

    # Arguments
        axes: a Matplotlib Axes object into which to plot
        image: an imshow-compatible Image or array to plot
        title: title to add to the plot
    """
    if (len(image.shape) == 2) or image.shape[-1] == 1:
        axes.imshow(image, cmap="Greys")
    else:
        axes.imshow(image)

    axes.set_xticklabels([])
    axes.set_yticklabels([])
    axes.set_xticks([])
    axes.set_yticks([])
    axes.set_title(title)


def torch_data_to_image_grid(dataset, rows=5, cols=5, shuffle=False, rng=local_rng):
    """
    Extract images from a torchvision dataset (with ToTensor transform)
    and grid them an imshow-compatible array.
    """
    count = rows * cols
    ix = rng.permutation(len(dataset))[:count] if shuffle else range(count)

    images = [np.moveaxis(dataset[ii][0].numpy(), 0, 2) for ii in ix]

    return np.concatenate(
        [
            np.concatenate(images[start : (start + cols)], axis=1)
            for start in range(0, count, cols)
        ],
        axis=0,
    )