utils.py

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# %% For external qt plot
def animate_scalar_qt(scalar,
                   pause_time,
                   colormap = 'magma',
                   c_range = 'global',
                   c_min = None, c_max = None):
    
    assert c_range in {'global','local','custom'}, f'Invalid input: {c_range}. Must be one either \'global\',\'local\',\'custom\'.'

    T = scalar.shape[2]

    if c_range == 'global':
        S_min = scalar.min()
        S_max = scalar.max()
    elif c_range == 'custom':
        assert c_min is not None and c_max is not None, 'Colour limits must be defined'
        S_min = c_min
        S_max = c_max

    fig, ax = plt.subplots()
    img = ax.imshow(scalar[:, :, 0], cmap=colormap)
    
    plt.axis('off') # Axis or no axis
    
    for i in range(T):
        S_tmp = scalar[:, :, i]  
        img.set_array(S_tmp)
        
        # Set colour range
        if c_range == 'custom':
            S_min = S_tmp.min()
            S_max = S_tmp.max()
        img.set_clim(S_min, S_max)
        
        fig.canvas.draw()
        fig.canvas.flush_events()
        
        plt.pause(pause_time)
        
# %% Inline viewing for Spyder
def animate_scalar_inline(scalar,
                   pause_time,
                   colormap = 'magma',
                   c_range = 'global',
                   c_min = None, c_max = None):
        
    assert c_range in {'global','local','custom'}, f'Invalid input: {c_range}. Must be one either \'global\',\'local\',\'custom\'.'
    
    T = scalar.shape[2]

    for i in range(T):
        S_tmp = scalar[:, :, i]  

        if c_range == 'global':
            S_min = scalar.min()
            S_max = scalar.max()
        elif c_range == 'local':
            S_min = S_tmp.min()
            S_max = S_tmp.max()
        elif c_range == 'custom':
            assert c_min is not None and c_max is not None, 'Colour limits must be defined'
            S_min = c_min
            S_max = c_max
        
        plt.imshow(S_tmp, cmap = colormap, vmin = S_min, vmax = S_max)
        plt.colorbar()
        plt.axis('off')
        plt.show()
                    
        plt.pause(pause_time)
        
# %% Saving a gif
def save_gif(scalar, save_name =  None, colormap = 'magma', c_range = 'global', c_min = None, c_max = None):
    # [!] Add duration parameter
    
    assert save_name is not None, 'Give a filename.'
    assert c_range in {'global','local','custom'}, f'Invalid input: {c_range}. Must be one either \'global\',\'local\',\'custom\'.'

    if c_min is None and c_max is None:
        S_min = scalar.min()
        S_max = scalar.max()
    if c_range == 'global': # Min and max of full array
        S_min = scalar.min()
        S_max = scalar.max()
    elif c_range == 'custom': # Custom colour limits
        assert c_min is not None and c_max is not None, 'Colour limits must be defined'
        S_min = c_min
        S_max = c_max
    
    print('Saving animation result...')
    T = scalar.shape[2]
    color_vals = plt.get_cmap(colormap)

    frames = []
    for i in range(T):
        S_tmp = scalar[:, :, i]  
        
        if c_range == 'local': # Colour limits of each individual frame
            S_min = S_tmp.min()
            S_max = S_tmp.max()

        S_norm = (S_tmp - S_min) / (S_max - S_min)
        colored_field = color_vals(S_norm)
        colored_field_rgb = (colored_field[:, :, :3] * 255).astype('uint8')
        tmp_image = Image.fromarray(colored_field_rgb)
        frames.append(tmp_image)
        
    frames[0].save(f'{save_name}.gif', save_all=True, append_images=frames[1:], duration=20, loop=0)
    print(f'Saved as: {save_name}.gif')
    
    
# %%
def plot_vec_field_on_image(X_in,Y_in,U_in,V_in,S_in,
                            DS = 1,
                            colormap = 'magma',
                            vec_color = 'black',
                            vec_scale = 1,
                            vec_alpha = 1,
                            min_S = None, max_S = None):
    
    # Downsample
    X_plot = X_in[::DS,::DS]
    Y_plot = Y_in[::DS,::DS]
    U_plot = U_in[::DS,::DS]
    V_plot = V_in[::DS,::DS]

    # Instantaneous min max limits if none specified
    if min_S is None and max_S is None:
        min_S = S_in.min()
        max_S = S_in.max()
    
    # Display
    plt.figure(figsize=(6,6))
    # Image
    plt.imshow(S_in, cmap = colormap, vmin = min_S, vmax = max_S)
    plt.colorbar(shrink=0.7)
    # Vector field
    plt.quiver(X_plot, Y_plot, U_plot, V_plot, 
               color = vec_color, 
               scale = vec_scale, 
               alpha = vec_alpha,
               angles = 'xy')
    plt.axis('off')
    plt.show()

# %%

def plot_vec_field(X_in,Y_in,U_in,V_in,
                                DS = 1,
                                vec_color = 'black',
                                vec_scale = 1,
                                vec_alpha = 1):
    
    # Downsample
    X_plot = X_in[::DS,::DS]
    Y_plot = Y_in[::DS,::DS]
    U_plot = U_in[::DS,::DS]
    V_plot = V_in[::DS,::DS]

    # Display
    plt.figure(figsize=(6,6))
    # Vector field
    plt.quiver(X_plot, Y_plot, U_plot, V_plot, 
               color = vec_color, 
               scale = vec_scale, 
               alpha = vec_alpha,
               angles = 'xy')
    plt.axis('off')
    plt.show()