python matplotlib Line Collection change color of some lines

Question

I have to represent deformations of a grid of size 256x256, from what I did with this functions

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection


def plot_grid(x,y, ax=None, **kwargs):
    ax = ax or plt.gca()
    segs1 = np.stack((x,y), axis=2)
    segs2 = segs1.transpose(1,0,2)
    ax.add_collection(LineCollection(segs1, **kwargs))
    ax.add_collection(LineCollection(segs2, **kwargs))

    ax.autoscale()

def plot_diff(diff,directory,k,color_down,color_up,name):
    u, v = diff[:,:,0], diff[:,:,1]

    nanIdx = np.logical_or(np.isnan(u), np.isnan(v))
    u[np.where(nanIdx)] = 0
    v[np.where(nanIdx)] = 0
    
    plt.figure(figsize=(50,50))
    
    plot_grid(u, v, color=color_down, linewidth=0.5)
   
    plt.show()
    plt.savefig(directory + "/" + name + "-" + str(k) + ".png")
    plt.close()

calling plot_diff on the np.array of size (256,256,2) describing the deformation I obtain this kind of representation:

which is kind of hard to look at for the density of all the lines. So what I would have liked to do is being able to select only some lines (for example every 8 lines) and plot them in a different color and with a stronger linewidth, so that the end result is more pleasant to look at.

The only thing that I was able to do was selecting a 'subgrid' from the deformation (deleting data from u and v) and plot that on top of the other, with the following code

def axes_to_delete(size):
    res = np.arange(0,size)
    return np.delete(res,np.arange(0,size,8))[:-1]   

def less_lines(x,y):
    size = x.shape[0]
    delete = axes_to_delete(size)
    return np.delete(np.delete(x,delete,0),delete,1), np.delete(np.delete(y,delete,0),delete,1)
        
def plot_diff(diff,directory,k,color_down,color_up,name,max_flow=-1,min_max_flow=-1):
    u, v = diff[:,:,0], diff[:,:,1]
    nanIdx = np.logical_or(np.isnan(u), np.isnan(v))
    u[np.where(nanIdx)] = 0
    v[np.where(nanIdx)] = 0
    
    plt.figure(figsize=(50,50))
  
    plot_grid(u, v, color=color_down, linewidth=0.5)
    uu,vv = less_lines(u,v)
   
    plot_grid(uu, vv, color=color_up, linewidth=2.5)
    plt.show()
    plt.savefig(directory + "/" + name + "-" + str(k) + ".png")
    plt.close()

and in this way I obtain the following plot which is a bit nicer to look at, but the problem is that of course these new red lines are linearly interpolated, so I have a coarser red grid on top of the grey one (it's visible zooming in), and my goal is to have the exact lines as below only in a different color and width, without loosing the data from the original lines.

The problem is that I don't know LineCollection that well and I couldn't find a way to do that. Does anyone have any idea of what could do it?

Answer 1

Acording to the official docs for matplotlib.collections.LineCollection, the normal plot parameters used in plot functions like matplotlib.pyplot.plot can be passed as a single element OR a list that will be cycled through for each line. Use the colors parameter to make every nth line a different color.

For example, to make all lines except every 4th red: colors=((0,0,0),(0,0,0),(0,0,0),(1,0,0))