Source code for pyparadigm.extras

"""Using functionality from this file requires extra dependencies."""
import warnings

try:
    import numpy as np
    from matplotlib import cm
    import matplotlib as mpl
except ImportError:
    warnings.warn("This module requires numpy and matplotlib, which seem to be missing.")

import contextlib
with contextlib.redirect_stdout(None):
    import pygame


def _normalize(mat: np.ndarray):
    """rescales a numpy array, so that min is 0 and max is 255"""
    abs_max = np.nanmax(np.abs(mat))
    norm = mpl.colors.Normalize(vmin=-float(abs_max),vmax=abs_max)
    return norm(mat).data


[docs]def to_24bit_gray(mat: np.ndarray): """returns a matrix that contains RGB channels, and colors scaled from 0 to 255""" return np.repeat(np.expand_dims(_normalize(mat) * 255, axis=2), 3, axis=2).astype(np.uint8)
[docs]def apply_color_map(name: str, mat: np.ndarray = None, normalize = True): """returns an RGB matrix scaled by a matplotlib color map, if normalize is False the matrix must only have values in the range [0, 1]""" def apply_map(mat): if normalize: mat = _normalize(mat) return (cm.get_cmap(name)(mat) * 255).astype(np.uint8) return apply_map if mat is None else apply_map(mat)
[docs]def mat_to_surface(mat: np.ndarray, transformer=to_24bit_gray): """Can be used to create a pygame.Surface from a 2d numpy array. By default a grey image with scaled colors is returned, but using the transformer argument any transformation can be used. :param mat: the matrix to create the surface of. :type mat: np.ndarray :param transformer: function that transforms the matrix to a valid color matrix, i.e. it must have 3dimension, were the 3rd dimension are the color channels. For each channel a value between 0 and 255 is allowed :type transformer: Callable[np.ndarray[np.ndarray]]""" transformed_mat = transformer(mat.T) if transformer is not None else mat.T if transformed_mat.shape[2] == 4: img = pygame.pixelcopy.make_surface(transformed_mat[:, :, :3])\ .convert_alpha() pix_arr = pygame.surfarray.pixels_alpha(img) pix_arr = transformed_mat[:, :, 3] return img return pygame.pixelcopy.make_surface(transformed_mat)