Python Gravity Field Calculations

Question

I am relatively new to programming, only a few months at university. I am trying to create a gravity field strength visualization pygame program. For that I have a function which calculates the gravity field at each pixel due to a certain number of masses randomly distributed around the given screen. This function is a modified version of my gravitational force calculator, it calculates the total gravitational force experienced by a mass due to a certain number of masses and it works as expected. The field calculator function is similar but slightly modified, although, when I print the 2D array, which contains the same amount of elements as pixels in the screen of pygame, all elements are 0 apart from the bottom row and I cannot seem to figure out why.

#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
def gravfield(mass, x, y, screen_width, screen_height):

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

    #Define array containing all field magnitudes at each pixel position of the screen
    field = np.zeros((screen_width, screen_height))   
        
    #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 lists which contains the field contributions of each mass
            x_fields = []
            y_fields = []
            
            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]

                #Calculate distance between the mass and point and also provide a unique value to the pixel which contains the ball to prevent division by zero
                distance = np.sqrt((x_distance**2) + (y_distance**2))
                if distance == 0:
                    x_fields = 0
                    y_fields = 0
                    
                else:
                    

                        #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 a list which contains field components caused by all the other masses
                        x_field = (magnitude * x_unit_vector) 
                        y_field = (magnitude * y_unit_vector) 
                        x_fields.append(x_field)
                        y_fields.append(y_field)

                
        #Sum all the field components in each direction to get the resultant field and then its magnitude which is assigned to an array where each position corresponds to a pixel on the screen
           
        gx = sum(x_fields)
        gy = sum(y_fields)
        g = np.sqrt((gx**2) + (gy**2))
        field[n, i] = g

    return field

This is my current field function, it takes a list of masses, their coordinates and the dimensions of the screen and returns a field array from numpy where each element corresponds to a pixel on the screen. I expected the values to not be all zeros almost everywhere. I tried using a screen of 100x100 with one mass at (50, 50) with a mass of 6,000,000,000,000 but still nothing.

Answer 1

You need to vectorise, which basically means deleting all of your code. Numpy is not intended to run in Python loops (the loops are written for you in C).

The choice of pygame is questionable, depending on what you're doing; I demonstrate with matplotlib instead.

Don't define G yourself; import it from scipy.

import matplotlib.pyplot as plt
import numpy as np
import scipy.constants
from matplotlib.colorbar import Colorbar
from matplotlib.image import AxesImage
from matplotlib.ticker import MultipleLocator


def grav_field(
    mass_kg: np.ndarray,
    y_m: np.ndarray,
    x_m: np.ndarray,
    screen_y_pixels: int = 600,
    screen_x_pixels: int = 800,
    m_per_pixel: float = 1e9,
) -> tuple[
    np.ndarray,  # y axis (m)
    np.ndarray,  # x axis (m)
    np.ndarray,  # field
    np.ndarray,  # magnitude
]:
    """
    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
    """

    # Example dimensions shown for 4 bodies, 800x600 screen
    # Assume that coordinate (0, 0) in orbital space is centred on the screen
    y = np.linspace(
        start=-screen_y_pixels*m_per_pixel/2,
        stop= +screen_y_pixels*m_per_pixel/2,
        num=screen_y_pixels,
    )[:, np.newaxis]  # 600x1
    x = np.linspace(
        start=-screen_x_pixels*m_per_pixel/2,
        stop= +screen_x_pixels*m_per_pixel/2,
        num=screen_x_pixels,
    )[np.newaxis, :]  # 1x800

    # Displacement components for each axis, position and body
    ydiff = y_m[np.newaxis, np.newaxis, :] - y[..., np.newaxis]  # 600x1x4
    xdiff = x_m[np.newaxis, np.newaxis, :] - x[..., np.newaxis]  # 1x800x4
    displacement = np.stack(np.broadcast_arrays(xdiff, ydiff))  # 2x600x800x4

    # Frobenius norm (rectilinear distance) over xy
    distance = np.linalg.norm(displacement, axis=0)  # 600x800x4

    # Displacement unit vectors for each axis, position and body
    unit = displacement / distance  # 2x600x800x4

    # Field magnitude for each position and body
    body_magnitudes = scipy.constants.G * mass_kg / distance**2  # 600x800x4

    # Resultant field for each axis and position
    field = (body_magnitudes*unit).sum(axis=-1)  # 2x600x800

    # Frobenius norm (magnitude for each position)
    magnitude = np.linalg.norm(field, axis=0)  # 600x800

    return y, x, field, magnitude


def plot(
    names: tuple[str, ...],
    y_body: np.ndarray,
    x_body: np.ndarray,
    y: np.ndarray,
    x: np.ndarray,
    field: np.ndarray,
    magnitude: np.ndarray,
) -> plt.Figure:
    fig: plt.Figure
    left: plt.Axes
    right: plt.Axes

    fig, (left, right) = plt.subplots(ncols=2)
    fig.suptitle('Gravitational field, after the sun disappeared')

    im: AxesImage = left.imshow(
        X=magnitude[::-1, :], norm='log',
        extent=(x[0], x[-1], y[0], y[-1]),
    )
    loc = MultipleLocator(1e11)
    left.yaxis.set_major_locator(loc)
    left.xaxis.set_major_locator(loc)
    cbar: Colorbar = fig.colorbar(mappable=im, ax=left)
    cbar.ax.set_ylabel('Field magnitude', rotation=-90)

    # streamplot is so slow as to be broken, so don't use it
    reduce = 25
    scale = np.sqrt(magnitude[::reduce, ::reduce])
    xx, yy = np.meshgrid(x[::reduce], y[::reduce])
    right.quiver(
        xx, yy,
        field[0, ::reduce, ::reduce] / scale,
        field[1, ::reduce, ::reduce] / scale,
        pivot='tip',
    )

    for name, yb, xb in zip(names, y_body, x_body):
        left.annotate(text=name, xy=(xb, yb), xytext=(xb + 1e10, yb - 3e10))
        right.scatter(xb, yb, c='red', marker='+')

    return fig


def test() -> None:
    names = 'Mercury', 'Venus', 'Earth', 'Mars'

    # Arbitrary orbital arguments
    arg = np.deg2rad((10, 25, 190, 250))

    # Vaguely correct radii, metres
    r = (50e9, 108e9, 150e9, 230e9)

    y_body = r*np.sin(arg)
    x_body = r*np.cos(arg)

    y, x, field, magnitude = grav_field(
        mass_kg=np.array((3.3011e23, 4.8675e24, 5.972168e24, 6.4171e23)),
        y_m=y_body, x_m=x_body,
    )

    fig = plot(names, y_body, x_body, y.ravel(), x.ravel(), field, magnitude)
    fig.tight_layout()
    plt.show()


if __name__ == '__main__':
    test()