Highlighting partitions in distinct color

Question

Given:

• G is an arbitrary graph
• partitions is a list of its edges divided into (in this case) 3 sets.

Color generation:

• Its purpose is to generate one distinct color for each partition.
• The function turns a number between 0 and 1 into an RGB color.
• For the given example, with 3 partitions, it would be called with 1/3, 2/3, and 1, and return a list of codes for green, blue, and red.
• For other numbers it would be used to spread the different colors as widely as possible.
• Visually, think of a circle with R, G, and B equally spaced around it. The number is the fraction of the way around the circle from R to pick the color, so if the number is .4, the color would be between G and B, closer to the G.

Color assignment:

• For each set in the partition, a different color is generated and assigned to a list of colors.
• In this case edge_color_list would be 9 greens, 4 blues, and 2 reds.
• The actual list is: [(0, 1.0, 0.0), (0.0, 0, 1.0), (0.0, 0, 1.0), (0, 1.0, 0.0), (0, 1.0, 0.0), (0, 1.0, 0.0), (0, 1.0, 0.0), (0, 1.0, 0.0), (0, 1.0, 0.0), (0, 1.0, 0.0), (0.0, 0, 1.0), (0, 1.0, 0.0), (1.0, 0, 0.0), (1.0, 0, 0.0), (0.0, 0, 1.0)]

Problem: The top and bottom sections can't really be changed, but the def_color() function and edge_color_list sections each look like they were written in C, and I think they could be done more elegantly.

It's obvious how I was thinking while writing it (i.e. how I would write it in C), but I'd like to know how python coders think.

I'm having trouble writing python without a heavy C accent (I wrote C for decades).

I understand the python language fairly well (or know how to recognize what I don't know and how to look it up).

And I've reached the stage where what I've written feels wrong.

But I still don't seem to think in native python:

#!/usr/bin/env python3
import matplotlib.pyplot as plot
import networkx as nx

G = nx.petersen_graph()
partitions = [{0, 3, 4, 5, 6, 7, 8, 9, 11}, {1, 2, 10, 14}, {12, 13}]

def get_color(fraction):
    if fraction < 1/3:
        color = (1-3*fraction, 3*fraction, 0)
    elif fraction < 2/3:
        fraction -= 1/3
        color = (0, 1-3*fraction, 3*fraction)
    else:
        fraction -= 2/3
        color = (3*fraction, 0, 1-3*fraction)
    return color

edge_color_list = [None] * len(G.edges())
for i in range(len(partitions)):
    items = list(partitions[i])   # convert from set
    for j in range(len(items)):
        edge_color_list[items[j]] = get_color((i+1)/len(partitions))

nx.draw(G, with_labels=True, pos=nx.circular_layout(G), width=2,
        node_color='pink', edge_color=edge_color_list)
plot.show()

Any suggestions about how I could have thought differently while writing the above would be appreciated.

Answer 1

I would change this

for i in range(len(partitions)):
    items = list(partitions[i])   # convert from set
    for j in range(len(items)):
        edge_color_list[items[j]] = get_color((i+1)/len(partitions)

to:

for idx, partition in enumerate(partitions): # I personally do it your way 
    items = list(partition)   # convert from set
    for item in items: 
        edge_color_list[items] = get_color((idx+1)/len(partitions) 
        # idx is the `i` from before 

I referenced this.

There might be a linting plugin to help align your code towards being pythonic