pythonmatplotlibvectorgravity

Python Matplotlib Quiver Plotting Vector Field

I am trying to plot a gravitational vector field with Matplotlib's quiver function, I have a function which calculates the gravitational field caused by any amount of masses but when plotting it, I get strange results around the mass' position. The function returns, for the purpose of the vector plotting, field components in x and y direction, their x and y coordinates and a 2D field array which is not necessary for the plotting but for other code later on not relevant to this. The x and y components are lists where each values in the 4 lists for each index correspond to the same element, so at index 0 in all 4 lists we can obtain the x and y component of the field and the position of it.

The current code is as follows:

import matplotlib.pyplot as plt
import numpy as np
import math

G = 6.674 * 10**(-11)

#Function to calculate gravitational field at each point on the screen caused by all other masses whose mass and positions are defined in lists along with the size of the screen and the components of the field
def gravfield(mass, x, y, screen_width, screen_height, give_components = False):

    #Define screen position points
    x_screen = list(range(screen_width))
    y_screen = list(range(screen_height))

    #Define empty array containing all field magnitudes at each pixel position of the screen along with empty field component lists and positions
    field = np.zeros((screen_width, screen_height))
    field_x = []
    field_y = []
    field_x_position = []
    field_y_position = []
        
    #Calculate the gravitational field from each mass by the vector method    
    for i in range(len(x_screen)):
        for n in range(len(y_screen)):

            #Define empty variable which will end up being the value of the field at a specific position after all mass contributions have been added
            x_field = 0
            y_field = 0
            
            for j in range(len(mass)):
                
                #Calculate position vector from the mass causing the field to the point affected
                x_distance = x[j] - x_screen[i]
                y_distance = y[j] - y_screen[n]
                distance = np.sqrt((x_distance**2) + (y_distance**2))
                
                if not math.isclose(distance, 0):
                        #Calculate the unit vector of the shortest path between the mass and point
                        x_unit_vector = x_distance / distance
                        y_unit_vector = y_distance / distance

                        #Calculate magnitude of the field
                        magnitude = (G * mass[j] * 1) / (distance ** 2)

                        #Calculate the component of the field in each direction and add them to the variable which contains the resultant field value at that position
                        x_field = (magnitude * x_unit_vector) + x_field
                        y_field = (magnitude * y_unit_vector) + y_field
                else:
                    pass
                
        #Add the resultant field component value and its position to the defined lists and calculate its magnitude which is assigned to an array where each position corresponds to a pixel on the screen

            #Only fill the field component lists if specified 
            if give_components == True:
                gx = x_field
                gy = y_field
                field_x.append(gx)
                field_y.append(gy)
                g = np.sqrt((gx**2) + (gy**2))
                field[n, i] = g
                field_x_position.append(i)
                field_y_position.append(n)
                
            else:
                gx = x_field
                gy = y_field
                g = np.sqrt((gx**2) + (gy**2))
                field[n, i] = g

    #Return the specified elements        
    if give_components == True:
        
        return field, field_x, field_y, field_x_position, field_y_position
    else:
        return field

mass = [400000000000000000000000, 800000000000000000000000]
x_cord = [10, 15]
y_cord = [10, 15]
screen_width = 20
screen_height = 20
give_components = True
field, field_x, field_y, field_x_position, field_y_position = gravfield(mass, x_cord, y_cord, screen_width, screen_height, give_components)



fig, ax = plt.subplots()

# Plot the vectors using quiver
ax.quiver(field_x_position, field_y_position, field_x, field_y)

# Set the x and y axis limits
ax.set_xlim([min(field_x_position), max(field_x_position)])
ax.set_ylim([min(field_y_position), max(field_y_position)])

# Show the plot
plt.show()

This function works when representing the intensity of the field in a color-coded form, but when plotting it with the quiver function I obtain what you can see in the image. The field is correct apart from near the mass itself where interesting vectors arise. I do not think it is a divide by zero error since those I just ignore when calculating.

Gravity field representation

Solution

  • I believe your solution is correct, but when using the default parameters (namely, angles="uv" rather than angles="xy") or not using ax.set_aspect(1) (equivalent: ax.set_aspect("equal")), the vectors can look wrong. After making one of those changes, the results look fine, as shown below. I also made some other edits to your code to simplify things (note: there is no need to do if value == True because you can just do if value).

    import matplotlib.pyplot as plt
import numpy as np
import math

plt.close("all")

G = 6.674e-11

# Function to calculate gravitational field at each point on the screen caused 
# by all other masses whose mass and positions are defined in lists along with 
# the size of the screen and the components of the field
def gravfield(mass, x, y, screen_width, screen_height, give_components=False):

    # Define screen position points
    x_screen = list(range(screen_width))
    y_screen = list(range(screen_height))

    # Define empty array containing all field magnitudes at each pixel position
    # of the screen along with empty field component lists and positions
    field = np.zeros((screen_width, screen_height))
    field_x = []
    field_y = []
    field_x_position = []
    field_y_position = []
        
    # Calculate the gravitational field from each mass by the vector method    
    for i in range(len(x_screen)):
        for n in range(len(y_screen)):

            # Define empty variable which will end up being the value of the 
            # field at a specific position after all mass contributions have 
            # been added
            x_field = 0
            y_field = 0
            
            for j in range(len(mass)):
                
                # Calculate position vector from the mass causing the field to 
                # the point affected
                x_distance = x[j] - x_screen[i]
                y_distance = y[j] - y_screen[n]
                distance = np.sqrt((x_distance**2) + (y_distance**2))
                
                if not math.isclose(distance, 0):
                    # Calculate the unit vector of the shortest path between 
                    # the mass and point
                    x_unit_vector = x_distance / distance
                    y_unit_vector = y_distance / distance
    
                    # Calculate magnitude of the field
                    magnitude = (G*mass[j]*1)/distance**2
    
                    # Calculate the component of the field in each 
                    # direction and add them to the variable which contains 
                    # the resultant field value at that position
                    x_field += magnitude*x_unit_vector
                    y_field += magnitude*y_unit_vector
                
            # Add the resultant field component value and its position to the 
            # defined lists and calculate its magnitude which is assigned to an 
            # array where each position corresponds to a pixel on the screen
            field[n, i] = np.sqrt((x_field**2) + (y_field**2))
    
            # Only fill the field component lists if specified 
            if give_components:
                field_x_position.append(i)
                field_y_position.append(n)
                field_x.append(x_field)
                field_y.append(y_field)

    # Return the specified elements        
    if give_components:
        return field, field_x, field_y, field_x_position, field_y_position
    return field

mass = [4000]
x_cord = [10]
y_cord = [10]
screen_width = 20
screen_height = 20
give_components = True
field, field_x, field_y, field_x_position, field_y_position = gravfield(mass, x_cord, y_cord, screen_width, screen_height, give_components)

field_x = np.array(field_x)
field_y = np.array(field_y)


fig, ax = plt.subplots()

# Plot the vectors using quiver
ax.quiver(field_x_position, field_y_position, field_x, field_y)
ax.set_aspect(1)

# Set the x and y axis limits
ax.set_xlim([min(field_x_position), max(field_x_position)])
ax.set_ylim([min(field_y_position), max(field_y_position)])

# Show the plot
fig.show()