Assign line a color by its angle in matplotlib-CodePudding

I'm looking for a way to assign color to line plots in matplotlib in a way that's responsive to the line's angle. This is my current code:

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
horz = [[0.5,0.6,0.8],[0.1,0.8,0.9],[0.2,0.5,0.9]] 
vert = [[0.1,0.2,0.3],[0.05,0.1,0.15],[0.2,0.3,0.35]]

f = plt.figure(figsize=(6,6))
ax = plt.axes()

for column in range(0,len(horz)):
    x = np.array(horz[column])
    y = np.array(vert[column])
            
    #LINEAR TRENDLINE
    z = np.polyfit(horz[column], vert[column], 1)
    p = np.poly1d(z)
    ax.plot(horz[column],p(horz[column]),"-")
    plt.arrow(x=horz[column][-2],y=p(horz[column])[-2],dx=(horz[column][-1]-horz[column][-2]),dy=(p(horz[column])[-1]-p(horz[column])[-2]), shape='full', lw=.01,
              length_includes_head=True, head_width=.012, head_length=0.02, head_starts_at_zero=False, overhang = 0.5)
            
    #FIG SETTINGS
    plt.xlim([0, 1])
    plt.ylim([0.1,0.5])
    ax.set_title('Title',
             fontsize = 14)

The idea here would be that if the line is at 0 degrees, it would be at one end of a given gradient, and if it were at 90 degrees, at the other end. Additionally, I'd like the line length to be taken as the intensity of the color. So if the line is short, it'd be closer to white, and if the line is long, it'd be closer to the raw color from the gradient.

CodePudding user response：

Managed to solve it myself. Used pretty simple formulas for calculating the lines' slopes and distances and then used these as input for the color mapping and alpha transparency attribute.

import geopandas as gpd
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.colors as colors
import numpy as np

%matplotlib inline

#Data
horz = [[0.5,0.6,0.8],[0.1,0.3,0.4],[0.2,0.5,0.9],[0.9,0.95,0.95]] 
vert = [[0.1,0.2,0.45],[0.05,0.1,0.15],[0.2,0.3,0.35],[0.1,0.3,0.5]]
            

#Slope calculation
def slopee(x1,y1,x2,y2):
        x = (y2 - y1) / (x2 - x1)
        return x

#Color set up
cmap = plt.cm.coolwarm_r
#0 means a horizontal line, 1 means a line at 45 degrees, infinite means a vertical line (2 is vertical enough)
cNorm  = colors.Normalize(vmin=0, vmax=2)
scalarMap = cm.ScalarMappable(norm=cNorm,cmap=cmap)

#Fig settings
f = plt.figure(figsize=(6,6))
ax = plt.axes()


for column in range(0,len(horz)):
    x = np.array(horz[column])
    y = np.array(vert[column])
    
    #LINEAR TRENDLINE
    # 1   LINEAR
    # >=2 POLINOMIAL
    z = np.polyfit(horz[column], vert[column], 1)
    p = np.poly1d(z)
    
    #Distance calc formula
    def calculateDistance(x1,y1,x2,y2):
        dist = np.sqrt((x2 - x1)**2   (y2 - y1)**2)
        return dist
    
    #Set up max an min distances
    maxdist = calculateDistance(0,0,0,0.9)
    mindist = calculateDistance(0,0,0,0)
    
    #Calculate line slope
    slope = slopee(horz[column][0],p(horz[column])[0],horz[column][-1],p(horz[column])[-1])
    
    #Not interested in any slopes going "down"
    if slope >=0:
        #Map colors based on slope (0-2)
        colorVal = scalarMap.to_rgba(slope)
        #Map transparency based on distance
        transparency = (calculateDistance(horz[column][0],p(horz[column])[0],horz[column][-1],p(horz[column])[-1])-mindist)/(maxdist-mindist)
        #Set up minimun transparency to be 50% instead of 0%
        transparency = (0.5*transparency)   0.5
        #The actual arrow plot
        plt.arrow(x=horz[column][0],y=p(horz[column])[0],dx=(horz[column][-1]-horz[column][0]),dy=(p(horz[column])[-1]-p(horz[column])[0]), shape='full',length_includes_head=True, head_starts_at_zero=False, lw=.5, head_width=.011, head_length=0.01, overhang = 0.5, color=colorVal,alpha=transparency)

    #FIG SETTINGS
    plt.xlim([0, 1])
    plt.ylim([0,0.5])
    ax.set_title('Title',fontsize = 14)

CodePudding user response：

Congrats on solving it yourself. I had put this together before I realized you had posted your answer. Very similar approach:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm, colors
from math import sqrt
plt.rcParams["figure.figsize"] = (15,15)

# Create a color mapper for degress to color
values = np.linspace(1.0, 90.0, 90)
norm = colors.Normalize(vmin=0.0, vmax=90.0, clip=True)
mapper = cm.ScalarMappable(norm=norm, cmap=cm.coolwarm_r)

horz = [[0.5,0.6,0.8],[0.1,0.3,0.4],[0.2,0.5,0.9],[0.9,0.95,0.95]] 
vert = [[0.1,0.2,0.45],[0.05,0.1,0.15],[0.2,0.3,0.35],[0.1,0.3,0.5]]
            

f = plt.figure(figsize=(15,15))
ax = plt.axes()

# Calculate lengths of each line
lengths =  [sqrt((x[-1]-x[0])**2   (y[-1]-y[0])**2) for x,y in zip(horz, vert)]

for x,y,length in zip(horz,vert, lengths):
    alpha = length / sum(lengths)
    angle = np.rad2deg(np.arctan2(y[-1] - y[0], x[-1] - x[0]))
    color = mapper.to_rgba(angle)    
    plt.arrow(x[0],y[0],x[-1]-x[0], y[-1]-y[0],shape='full', lw=2,
              length_includes_head=True, head_width=.012, head_length=0.02, head_starts_at_zero=False, overhang = 0.5, alpha=alpha, color=color)