import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

D = 3.0    # hranica x/y/z hodnot
R = 2.0    # polomer galaktickeho disku
d = 0.11    # suradnicove delenie
Sigma0 = 1
a = 1
x1 = 0
FI = []
for y1 in np.arange(-D, D, d):
    # Pomocou vhodneho d odstranime body z1 = 0, v ktorych je potencial -Inf
    for z1 in np.arange(-D, D, d):
        fi = 0
        for x2 in np.arange(-R, R, d):
            for y2 in np.arange(-R, R, d):
                r = np.sqrt(x2**2 + y2**2)
                # Vzdialenost bodu [x1, y1, z1] od jednotlivych elementov
                # disku [x2, y2 z2=0].
                if r <= R:
                    r2 = np.sqrt((x1 - x2)**2 + (y1 - y2)**2 + z1**2)
                    # Definicia plosnej hustoty pre bod o polohe [x2, y2].
                    Sigma = Sigma0*(1 + (r/a)**2)**(-3/2)
                    # Hmotnost elementu sa urci ako plosna hustota v pozicii
                    # [x2, y2, z2=0] nasobena plosnym diferencialom.
                    M = Sigma*d**2
                    # Celkovy potencial je dany suctom ciastkovych potencialov
                    # od jednotlivych bodov.
                    fi -= M/r2
        FI.append([y1, z1, fi])
FI = np.array(FI)
plt.figure()
plt.scatter(FI[:, 0], FI[:, 1], c=FI[:, 2], cmap=cm.coolwarm, lw = 0, s=10)
plt.colorbar()
plt.axis('equal')
plt.xlabel('y [a]')
plt.ylabel('z [a]')





