def test_calculate_trajectory_iterator(
pos_vec,
vel_vec,
time,
M,
start_lambda,
end_lambda,
OdeMethodKwargs,
return_cartesian,
):
cl1 = Schwarzschild.from_spherical(pos_vec, vel_vec, time, M)
arr1 = cl1.calculate_trajectory(
start_lambda=start_lambda,
end_lambda=end_lambda,
OdeMethodKwargs=OdeMethodKwargs,
return_cartesian=return_cartesian,
)[1]
cl2 = Schwarzschild.from_spherical(pos_vec, vel_vec, time, M)
it = cl2.calculate_trajectory_iterator(
start_lambda=start_lambda,
OdeMethodKwargs=OdeMethodKwargs,
return_cartesian=return_cartesian,
)
arr2_list = list()
for _, val in zip(range(100), it):
arr2_list.append(val[1])
arr2 = np.array(arr2_list)
assert_allclose(arr1[:100, :], arr2, rtol=1e-10)
def __get_x_y(self, coords, end_lambda, step_size):
"""
Parameters
----------
coords : ~einsteinpy.coordinates.velocity.SphericalDifferential
Initial position and velocity of particle in Spherical coordinates.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
"""
swc = Schwarzschild.from_spherical(coords, self.mass, self.time)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
# time = np.array([coord[0] for coord in vals])
r = np.array([coord[1] for coord in vals])
# theta = np.array([coord[2] for coord in vals])
phi = np.array([coord[3] for coord in vals])
x = r * np.cos(phi)
y = r * np.sin(phi)
self.__xarr = x
self.__yarr = y
return x, y
def plot(self, coords, end_lambda=10, step_size=1e-3):
"""
Parameters
----------
coords : ~einsteinpy.coordinates.velocity.SphericalDifferential
Position and velocity components of particle in Spherical Coordinates.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
"""
swc = Schwarzschild.from_spherical(coords, self.mass, self.time)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
plt.scatter(pos_x, pos_y, s=1, c=time, cmap=self.cmap_color)
if not self._attractor_present:
self._plot_attractor()
def plot(self, pos_vec, vel_vec, end_lambda=10, step_size=1e-3):
"""
Parameters
----------
pos_vec : list
list of r, theta & phi components along with ~astropy.units.
vel_vec : list
list of velocities of r, theta & phi components along with ~astropy.units.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
"""
swc = Schwarzschild.from_spherical(pos_vec, vel_vec, self.time,
self.mass)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
plt.scatter(pos_x, pos_y, s=1, c=time, cmap="Oranges")
if not self._attractor_present:
self._plot_attractor()
def test_calculate_trajectory2():
# based on the revolution of earth around sun
# data from https://en.wikipedia.org/wiki/Earth%27s_orbit
M = 1.989e30 * u.kg
distance_at_perihelion = 147.10e6 * u.km
speed_at_perihelion = 30.29 * u.km / u.s
angular_vel = (speed_at_perihelion / distance_at_perihelion) * u.rad
sph_obj = SphericalDifferential(
distance_at_perihelion,
np.pi / 2 * u.rad,
0 * u.rad,
0 * u.km / u.s,
0 * u.rad / u.s,
angular_vel,
)
end_lambda = ((1 * u.year).to(u.s)).value
cl = Schwarzschild.from_spherical(sph_obj, M)
ans = cl.calculate_trajectory(
start_lambda=0.0,
end_lambda=end_lambda,
OdeMethodKwargs={"stepsize": end_lambda / 2e3},
)[1]
# velocity should be 29.29 km/s at apehelion(where r is max)
i = np.argmax(ans[:, 1]) # index whre radial distance is max
v_apehelion = (((ans[i][1] * ans[i][7]) * (u.m / u.s)).to(u.km /
u.s)).value
assert_allclose(v_apehelion, 29.29, rtol=0.01)
def plot_trajectory(self, coords, end_lambda, step_size, color):
"""
Parameters
----------
coords : ~einsteinpy.coordinates.velocity.SphericalDifferential
Initial position and velocity of particle in Spherical coordinates.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
color : string
Color of the Geodesic
"""
swc = Schwarzschild.from_spherical(coords, self.mass, self.time)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
# time = np.array([coord[0] for coord in vals])
r = np.array([coord[1] for coord in vals])
# theta = np.array([coord[2] for coord in vals])
phi = np.array([coord[3] for coord in vals])
x = r * np.cos(phi)
y = r * np.sin(phi)
lines = self.ax.plot(x, y, "--", color=color)
return lines, x[-1], y[-1]
def animate(self, coords, end_lambda=10, step_size=1e-3, interval=50):
"""
Function to generate animated plots of geodesics.
Parameters
----------
coords : ~einsteinpy.coordinates.velocity.SphericalDifferential
Position and velocity components of particle in Spherical Coordinates.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
interval : int, optional
Control the time between frames. Add time in milliseconds.
"""
swc = Schwarzschild.from_spherical(coords, self.mass, self.time)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
frames = pos_x.shape[0]
x_max, x_min = max(pos_x), min(pos_x)
y_max, y_min = max(pos_y), min(pos_y)
margin_x = (x_max - x_min) * 0.1
margin_y = (y_max - y_min) * 0.1
fig = plt.figure()
plt.xlim(x_min - margin_x, x_max + margin_x)
plt.ylim(y_min - margin_y, y_max + margin_y)
pic = plt.scatter([], [], s=1, c=[])
plt.scatter(0, 0, color="black")
def _update(frame):
pic.set_offsets(
np.vstack((pos_x[:frame + 1], pos_y[:frame + 1])).T)
pic.set_array(time[:frame + 1])
return (pic, )
self.animated = FuncAnimation(fig,
_update,
frames=frames,
interval=interval)
def test_calculate_trajectory_iterator_RuntimeWarning2():
pos_vec = [306 * u.m, np.pi / 2 * u.rad, np.pi / 3 * u.rad]
vel_vec = [0 * u.m / u.s, 0.01 * u.rad / u.s, 10 * u.rad / u.s]
time = 0 * u.s
M = 1e25 * u.kg
start_lambda = 0.0
OdeMethodKwargs = {"stepsize": 0.4e-6}
cl = Schwarzschild.from_spherical(pos_vec, vel_vec, time, M)
with warnings.catch_warnings(record=True) as w:
it = cl.calculate_trajectory_iterator(
start_lambda=start_lambda,
OdeMethodKwargs=OdeMethodKwargs,
stop_on_singularity=False,
)
for _, _ in zip(range(1000), it):
pass
assert len(w) >= 1
def test_calculate_trajectory(coords, time, M, start_lambda, end_lambda,
OdeMethodKwargs):
cl = Schwarzschild.from_spherical(coords, M, time)
ans = cl.calculate_trajectory(
start_lambda=start_lambda,
end_lambda=end_lambda,
OdeMethodKwargs=OdeMethodKwargs,
)
ans = ans[1]
testarray = list()
for i in ans:
g = schwarzschild_utils.metric(i[1], i[2], M.value)
testarray.append(g[0][0] * (i[4]**2) + g[1][1] * (i[5]**2) + g[2][2] *
(i[6]**2) + g[3][3] * (i[7]**2))
testarray = np.array(testarray, dtype=float)
comparearray = np.ones(shape=ans[:, 4].shape, dtype=float)
assert_allclose(testarray, comparearray, 1e-4)
def plot_trajectory(self, pos_vec, vel_vec, end_lambda, step_size, color):
swc = Schwarzschild.from_spherical(pos_vec, vel_vec, self.time, self.mass)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size}
)[1]
time = np.array([coord[0] for coord in vals])
r = np.array([coord[1] for coord in vals])
theta = np.array([coord[2] for coord in vals])
phi = np.array([coord[3] for coord in vals])
x = r * np.cos(phi)
y = r * np.sin(phi)
lines = self.ax.plot(x, y, "--", color=color)
return lines, x[-1], y[-1]
def test_calculate_trajectory(pos_vec, vel_vec, time, M, start_lambda,
end_lambda, OdeMethodKwargs):
cl = Schwarzschild.from_spherical(pos_vec, vel_vec, time, M)
ans = cl.calculate_trajectory(
start_lambda=start_lambda,
end_lambda=end_lambda,
OdeMethodKwargs=OdeMethodKwargs,
)
_c, _scr = constant.c.value, schwarzschild_radius(M).value
ans = ans[1]
testarray = list()
for i in ans:
g = schwarzschild_utils.metric(_c, i[1], i[2], _scr)
testarray.append(g[0][0] * (i[4]**2) + g[1][1] * (i[5]**2) + g[2][2] *
(i[6]**2) + g[3][3] * (i[7]**2))
testarray = np.array(testarray, dtype=float)
comparearray = np.ones(shape=ans[:, 4].shape, dtype=float)
assert_allclose(testarray, comparearray, 1e-4)
def test_calculate_trajectory(
pos_vec, vel_vec, time, M, start_lambda, end_lambda, OdeMethodKwargs
):
cl = Schwarzschild.from_spherical(pos_vec, vel_vec, time, M)
ans = cl.calculate_trajectory(
start_lambda=start_lambda,
end_lambda=end_lambda,
OdeMethodKwargs=OdeMethodKwargs,
)
_c, _scr = constant.c.value, schwarzschild_radius(M).value
ans = ans[1]
testarray = (
(1 - (_scr / ans[:, 1])) * np.square(ans[:, 4])
- (np.square(ans[:, 5])) / ((1 - (_scr / ans[:, 1])) * (_c ** 2))
- np.square(ans[:, 1] / _c)
* (np.square(ans[:, 6]) + np.square(np.sin(ans[:, 2])) * np.square(ans[:, 7]))
)
comparearray = np.ones(shape=ans[:, 4].shape, dtype=float)
assert_allclose(testarray, comparearray, 1e-4)
def plot(self, coords, end_lambda=10, step_size=1e-3):
"""
Parameters
----------
coords : ~einsteinpy.coordinates.velocity.SphericalDifferential
Position and velocity components of particle in Spherical Coordinates.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
"""
swc = Schwarzschild.from_spherical(coords, self.mass, self.time)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
self.data.append(
go.Scatter(x=pos_x,
y=pos_y,
mode='markers',
name='plot_points',
marker=dict(size=5,
color=self.cmap_color,
line=dict(width=2, ))))
init_notebook_mode(connected=self.connected)
self.fig = dict(data=self.data, layout=self.layout)
# plt.scatter(pos_x, pos_y, s=1, c=time, cmap=self.cmap_color)
if not self._attractor_present:
self._plot_attractor()
def plot(self, pos_vec, vel_vec, end_lambda=10, step_size=1e-3):
swc = Schwarzschild.from_spherical(pos_vec, vel_vec, self.time, self.mass)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size}
)[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
plt.scatter(pos_x, pos_y, s=1, c=time, cmap="Oranges")
if not self._attractor_present:
self._plot_attractor()
def test_calculate_trajectory_iterator_RuntimeWarning():
sph_obj = SphericalDifferential(
306 * u.m,
np.pi / 2 * u.rad,
np.pi / 2 * u.rad,
0 * u.m / u.s,
0.01 * u.rad / u.s,
10 * u.rad / u.s,
)
M = 1e25 * u.kg
start_lambda = 0.0
OdeMethodKwargs = {"stepsize": 0.4e-6}
cl = Schwarzschild.from_spherical(sph_obj, M)
with warnings.catch_warnings(record=True) as w:
it = cl.calculate_trajectory_iterator(
start_lambda=start_lambda,
OdeMethodKwargs=OdeMethodKwargs,
stop_on_singularity=True,
)
for _, _ in zip(range(1000), it):
pass
assert len(w) >= 1
def animate(self, coords, end_lambda=10, step_size=1e-3, interval=50):
"""
Function to generate animated plots of geodesics.
Parameters
----------
coords : ~einsteinpy.coordinates.velocity.SphericalDifferential
Position and velocity components of particle in Spherical Coordinates.
end_lambda : float, optional
Lambda where iteartions will stop.
step_size : float, optional
Step size for the ODE.
interval : int, optional
Control the time between frames. Add time in milliseconds.
"""
swc = Schwarzschild.from_spherical(coords, self.mass, self.time)
vals = swc.calculate_trajectory(
end_lambda=end_lambda, OdeMethodKwargs={"stepsize": step_size})[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
x_max, x_min = max(pos_x), min(pos_x)
y_max, y_min = max(pos_y), min(pos_y)
margin_x = (x_max - x_min) * 0.1
margin_y = (y_max - y_min) * 0.1
self.data = [
dict(visible=False,
line=dict(color='#00CED1', width=6),
name='',
x=pos_x[:i],
y=pos_y[:i]) for i in range(0, len(pos_x), 100)
]
steps = []
for i in range(len(self.data)):
step = dict(
method='restyle',
args=['visible', [False] * len(self.data)],
)
step['args'][1][i] = True # Toggle i'th trace to "visible"
steps.append(step)
self.layout = dict(
sliders=[
dict(active=100,
currentvalue={"prefix": "Frequency: "},
pad={"t": 50},
steps=steps)
],
title='[EINSTEINPY] Plotly Animation',
hovermode='closest',
xaxis=dict(tickmode='linear',
ticks='outside',
tick0=int(x_min - margin_x),
dtick=int((x_max + margin_x - x_min - margin_x)) / 10,
ticklen=8,
tickwidth=4,
tickcolor='#000'),
yaxis=dict(tickmode='linear',
ticks='outside',
tick0=int(y_min - margin_y),
dtick=int((y_max + margin_y - y_min - margin_y) / 10),
ticklen=8,
tickwidth=4,
tickcolor='#000'),
updatemenus=[{
'type':
'buttons',
'buttons': [{
'label': 'Play',
'method': 'animate',
'args': [None]
}]
}])
self.fig = dict(data=self.data, layout=self.layout)
vel = [0 * u.m / u.s, 0 * u.rad / u.s, 1900 * u.rad / u.s]
# In[ ]:
vel, pos
# In[ ]:
get_ipython().system('pip install einsteinpy==0.1')
import einsteinpy
M = 6e24 * u.kg
pos = [130 * u.m, np.pi / 2 * u.rad, -np.pi / 8 * u.rad]
vel = [0 * u.m / u.s, 0 * u.rad / u.s, 1900 * u.rad / u.s]
swc = Schwarzschild.from_spherical(pos, vel, 0 * u.s, M)
vals = swc.calculate_trajectory(end_lambda=0.002,
OdeMethodKwargs={"stepsize": 5e-8})[1]
time = vals[:, 0]
r = vals[:, 1]
# Currently not being used (might be useful in future)
# theta = vals[:, 2]
phi = vals[:, 3]
pos_x = r * np.cos(phi)
pos_y = r * np.sin(phi)
import plotly.plotly as py
import plotly.graph_objs as go
