Python Gravity Field Visualization

Question

I am trying to create function in python which receives an array which has the value of the gravitational field at each pixel of a pygame screen and then assigns a color from a certain range to each pixel depending on the magnitude of the field at each pixel. Although, I am not managing to correctly assign color values to the pixel positions with my conditions.

I have made a function, which takes this field numpy array and the screen dimensions as an input. It generates a list of colors from red to yellow and also separates the range of field magnitudes into the same amount of section as the number of colors in order to assign a color to each range. It then goes over all values in the field array to check if they fall within any of these ranges, which they should fall into at least one, and assigns the corresponding color.

def fieldview(field, screen_width, screen_height):

    #Define color gradient list
    iterator = 0
    colors = []
    for i in range(52):   #52 colors
        color = (255, iterator, 0)
        colors.append(color)
        iterator = iterator + 5
    
    #Divide the range of field values into sections
    

    min_field_index = np.argmin(field)
    max_field_index = np.argmax(field)
    real_min_field_index = np.unravel_index(min_field_index, field.shape)
    real_max_field_index = np.unravel_index(max_field_index, field.shape)
    min_field_value = field[real_min_field_index]
    max_field_value = field[real_max_field_index]

    step = (max_field_value - min_field_value) / (52)
   
    
    field_levels = []
    field_value = min_field_value
    for i in range(52):
        field_level = (field_value, field_value + step)
        field_levels.append(field_level)
        field_value = field_value + step

    field_levels.reverse()
    
    #Assign a color index to each pixel according to the field value
    color_coordinate_values = np.zeros((screen_width, screen_height), dtype=np.object_)
    
    for i in range(52):
        condition = (field >= field_levels[i][0]) & (field <= field_levels[i][1])
        color = colors.index(colors[i])
        color_coordinate_values[condition] = color  
    

    return color_coordinate_values

But, in the end, the colors are not correctly assigned and I cannot figure out why the conditions don't work when I tried it on other arrays before. If you have ideas for how I can assign these color values in a different way to the field values to in the end represent their magnitude visually then I am also open to those ideas. Thank you.

Answer 1

Consider a much simpler way of doing this, like the following:

def field_view(field):
    min_ = field.min()
    max_ = field.max()
    green = (255 * (max_ - field)) // (max_ - min_)
    return np.pad(green[:, :, None],  # pad the green values with 255 and 0
                  ((0, 0), (0, 0), (1, 1)),
                  constant_values=((255, 0)))  

print(field_view(np.array([[0,1,2], [3,4,5], [6,9,7]])))

Alternatively:

def field_view(field):
    def const(n):
        return np.full_like(field, n)

    return np.transpose([
        const(255),
        (255 * (field.max() - field)) // field.ptp(),
        const(0)
    ], (1,2,0))