├── .flake8
├── .github
└── workflows
│ └── ci.yml
├── .gitignore
├── .vscode
├── launch.json
├── settings.json
└── tasks.json
├── LICENSE
├── README.md
├── _config.yml
├── doc
├── earth_black_hole_animation.png
├── earth_black_hole_animation_zoomed.png
├── earth_black_hole_plot.png
├── gros_logo.png
├── gros_logo.svg
├── gros_logo_1280x640.png
└── mercury_plot.png
├── setup.py
└── src
└── gros
├── __init__.py
├── examples
├── earth_black_hole_trajectory_animate.py
├── earth_black_hole_trajectory_plot_only.py
└── mercury_trajectory.py
├── metric
├── __init__.py
└── schwarzschild.py
├── tests
├── metric
│ └── test_schwarzschild.py
└── utils
│ └── test_datahandling.py
└── utils
├── __init__.py
├── const.py
├── datahandling.py
├── log.py
└── transforms.py
/.flake8:
--------------------------------------------------------------------------------
1 | [flake8]
2 | max-line-length = 120
3 |
--------------------------------------------------------------------------------
/.github/workflows/ci.yml:
--------------------------------------------------------------------------------
1 | name: Tests
2 |
3 | on:
4 | push:
5 | branches: [ "master" ]
6 | pull_request:
7 | branches: [ "master" ]
8 |
9 | jobs:
10 | build:
11 |
12 | runs-on: ubuntu-latest
13 | strategy:
14 | fail-fast: false
15 | matrix:
16 | python-version: ["3.9", "3.10", "3.11"]
17 |
18 | steps:
19 | - uses: actions/checkout@v4
20 | - name: Set up Python ${{ matrix.python-version }}
21 | uses: actions/setup-python@v3
22 | with:
23 | python-version: ${{ matrix.python-version }}
24 | - name: Install dependencies
25 | run: |
26 | python -m pip install --upgrade pip
27 | python -m pip install flake8 pytest
28 | # if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
29 | pip install .
30 | - name: Lint with flake8
31 | run: |
32 | flake8 . --count --exit-zero --max-complexity=10 --max-line-length=120 --statistics
33 | - name: Test with pytest
34 | run: |
35 | pytest -vvv src/gros/tests
36 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | /.venv/
2 | *.cpython-37.pyc
3 |
--------------------------------------------------------------------------------
/.vscode/launch.json:
--------------------------------------------------------------------------------
1 | {
2 | "version": "0.2.0",
3 | "configurations": [
4 | {
5 | "name": "Python: Current File",
6 | "type": "python",
7 | "request": "launch",
8 | "program": "${file}",
9 | "console": "integratedTerminal",
10 | "envFile": "${workspaceFolder}/.env",
11 | "env": {
12 | "PYTHONPATH": "${workspaceFolder}/src/:${workspaceFolder}/src/gros/:${env:PYTHONPATH}"
13 | }
14 | }
15 | ]
16 | }
17 |
--------------------------------------------------------------------------------
/.vscode/settings.json:
--------------------------------------------------------------------------------
1 | {
2 | "python.linting.pylintEnabled": false,
3 | "python.linting.flake8Enabled": true,
4 | "python.linting.enabled": true,
5 | "python.formatting.provider": "black",
6 | "python.pythonPath": ".venv/bin/python"
7 | }
--------------------------------------------------------------------------------
/.vscode/tasks.json:
--------------------------------------------------------------------------------
1 | {
2 | "version": "2.0.0",
3 | "tasks": [
4 | {
5 | "label": "echo workspaceFolder",
6 | "type": "shell",
7 | "command": "echo ${workspaceFolder}"
8 | }
9 | ]
10 | }
11 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2020 BjoB (Björn Barschtipan)
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 |
4 |
5 | -----------------
6 |
7 | # Status
8 |
9 |
10 | [](https://github.com/bytebat/gros/actions/workflows/ci.yml/badge.svg)
11 | [](https://img.shields.io/badge/powered--by-astropy-blue)
12 | [](https://img.shields.io/badge/powered--by-plotly-blue)
13 |
14 | # Overview
15 |
16 | **gros** is a python package to numerically calculate and simulate particle trajectories based on the field equations of general relativity. A user needs to define a certain metric by providing the mass of a central gravitational attractor and the start coordinates and velocity of the test particle.
17 |
18 | # Installation
19 |
20 | Clone the repository and install the package via pip in your choosen environment:
21 |
22 | ```sh
23 | pip install .
24 | ```
25 |
26 | # Theoretical background
27 |
28 | To simulate particle trajectories around a spherically symmetric body, we use the *Schwarzschild* solution of Einstein's field equations, describing the exterior spacetime for our use case. Starting from the field equations
29 |
30 | $$\large R_{\mu\nu}-\frac{1}{2}g_{\mu\nu}R=\frac{8\pi%20G}{c^4}T_{\mu\nu}$$
31 |
32 | with a vanishing energy momentum tensor
33 |
34 | $$\large T_{\mu\nu}=0$$
35 |
36 | the *Schwarzschild* metric can be derived as
37 |
38 | $$\large ds^2=g_{\mu\nu}dx^\mu dx^\nu=c^2 (1-\frac{r_s}{r})dt^2-\frac{1}{1-r_s/r}dr^2-r^2d\theta^2-r^2sin^2\theta d\phi^2$$
39 |
40 | with the *Schwarzschild radius* $r_s=2GM/c^2$ and the space time coordinates based on spherical coordinates $(x^0,x^1,x^2,x^3) \mapsto (ct, r,\theta, \phi)$.
41 |
42 | The intrinsic space time curvature can be derived from the metric by evaluating the *Christoffel symbols* given with
43 |
44 | $$\large \Gamma_{\alpha\nu}^{\beta}=\frac{1}{2}g^{\mu\beta}(\partial_{\alpha}g_{\mu\nu}+\partial_{\nu}g_{\mu\alpha}-\partial_\mu g_{\alpha\nu})$$
45 |
46 | After calculating these coefficients and using the proper time as parameter, the motion of of a particle in the gravitational field can be retrieved by solving the system of differential equations given with the geodesic equations
47 |
48 | $$\large \frac{d^2 x^\mu}{d\tau^2}+\Gamma_{\alpha\beta}^{\mu}\frac{dx^\alpha}{d\tau}\frac{dx^\beta}{d\tau}=0$$
49 |
50 | # Examples
51 |
52 | Some simple simulations can be found in the [examples](https://github.com/BjoB/gros/tree/master/src/gros/examples) directory.
53 |
54 | ## Particle on Mercury's orbit
55 |
56 | This example simulates a particle orbiting the sun with some initial orbital parameters taken from
59 |
61 |
62 | ## Earth as a black hole
63 |
64 | What if earth was a black hole? The according example shows how a particle would act in short distance of 30m. Especially the perihelion precession is visualized as a direct effect of general relativity. Additionally the gravitational time dilation can be tracked along the animation frames with τ as the proper time of the particle. t is the calculated coordinate time, which can be seen as the measured proper time of a hypothetical observer positioned infinitely far away from the gravitational center.
65 |
66 |
67 |
69 |
--------------------------------------------------------------------------------
/_config.yml:
--------------------------------------------------------------------------------
1 | theme: jekyll-theme-slate
--------------------------------------------------------------------------------
/doc/earth_black_hole_animation.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/doc/earth_black_hole_animation.png
--------------------------------------------------------------------------------
/doc/earth_black_hole_animation_zoomed.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/doc/earth_black_hole_animation_zoomed.png
--------------------------------------------------------------------------------
/doc/earth_black_hole_plot.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/doc/earth_black_hole_plot.png
--------------------------------------------------------------------------------
/doc/gros_logo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/doc/gros_logo.png
--------------------------------------------------------------------------------
/doc/gros_logo.svg:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
23 |
25 |
28 |
37 |
84 |
114 |
115 |
117 | image/svg+xml
118 |
122 |
123 |
124 |
129 |
243 |
248 |
251 |
256 |
356 |
361 |
461 |
466 |
566 |
571 |
671 |
676 |
776 |
781 |
881 |
886 |
986 |
991 |
1091 |
1096 |
1196 |
1201 |
1280 |
1285 |
1396 |
1411 |
1412 |
1417 | G R O S eneral elativity rbit imulator
1474 |
1475 |
1476 |
1477 |
--------------------------------------------------------------------------------
/doc/gros_logo_1280x640.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/doc/gros_logo_1280x640.png
--------------------------------------------------------------------------------
/doc/mercury_plot.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/doc/mercury_plot.png
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | from setuptools import setup, find_packages
4 |
5 | with open("README.md", "r") as fh:
6 | long_description = fh.read()
7 |
8 | setup(
9 | name='gros',
10 | version='0.1.0',
11 | description='General Relativity Orbit Simulator',
12 | url='https://github.com/BjoB/gros',
13 | long_description=long_description,
14 | long_description_content_type="text/markdown",
15 | author='Björn Barschtipan',
16 | author_email='',
17 | license='MIT',
18 | packages=find_packages(where='src'),
19 | package_dir={'': 'src'},
20 | zip_safe=False,
21 | setup_requires=['wheel'],
22 | python_requires='>=3.5',
23 | install_requires=[
24 | 'astropy>=4.0',
25 | 'colorlog',
26 | 'numpy',
27 | 'pandas',
28 | 'plotly>=4.0',
29 | 'scipy',
30 | ],
31 | extras_require={
32 | 'dev': [
33 | 'black',
34 | 'pycodestyle',
35 | 'pyflakes',
36 | 'pytest',
37 | ],
38 | },
39 | )
40 |
--------------------------------------------------------------------------------
/src/gros/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/src/gros/__init__.py
--------------------------------------------------------------------------------
/src/gros/examples/earth_black_hole_trajectory_animate.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | from numpy import pi
4 |
5 | from gros.metric import schwarzschild as ss
6 | from gros.utils import const
7 |
8 |
9 | def earth_black_hole_trajectory():
10 | ss_metric = ss.SchwarzschildMetric(
11 | M=const.M_earth,
12 | initial_vec_pos=[30, pi/2, 0],
13 | initial_vec_v=[0, 0, 20000]
14 | )
15 |
16 | traj = ss_metric.calculate_trajectory(proptime_end=1.2e-4, step_size=1e-6)
17 | traj.plot(animation_step_size=1e-6)
18 |
19 |
20 | if __name__ == "__main__":
21 | earth_black_hole_trajectory()
22 |
--------------------------------------------------------------------------------
/src/gros/examples/earth_black_hole_trajectory_plot_only.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | from numpy import pi
4 |
5 | from gros.metric import schwarzschild as ss
6 | from gros.utils import const
7 |
8 |
9 | def earth_black_hole_trajectory():
10 | ss_metric = ss.SchwarzschildMetric(
11 | M=const.M_earth,
12 | initial_vec_pos=[100, pi/2, 0],
13 | initial_vec_v=[0, 0, 5000]
14 | )
15 |
16 | traj = ss_metric.calculate_trajectory(proptime_end=0.005, step_size=1e-5)
17 | traj.plot()
18 |
19 |
20 | if __name__ == "__main__":
21 | earth_black_hole_trajectory()
22 |
--------------------------------------------------------------------------------
/src/gros/examples/mercury_trajectory.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | from numpy import pi
4 |
5 | from gros.metric import schwarzschild as ss
6 | from gros.utils import const
7 |
8 |
9 | def mercury_trajectory():
10 |
11 | # init metric with mass of sun and mercury orbit parameters
12 | ss_metric = ss.SchwarzschildMetric(
13 | M=const.M_solar,
14 | initial_vec_pos=[57.9e6 * 1000, pi/2, 0],
15 | initial_vec_v=[0, 0, (47.4 / 57.9e6)]
16 | )
17 |
18 | radius_sun = 6.9634e8
19 |
20 | # calculate trajectory for one orbital period
21 | traj = ss_metric.calculate_trajectory(proptime_end=88 * 24 * 3600, step_size=3600)
22 |
23 | # animate with one frame per day
24 | traj.plot(attractor_radius=radius_sun, animation_step_size=24 * 3600)
25 |
26 | # comment out to show mercury's perihelion precession:
27 | # traj = ss_metric.calculate_trajectory(proptime_end=100 * 88 * 24 * 3600, step_size=86400)
28 | # traj.plot(attractor_radius=radius_sun)
29 |
30 |
31 | if __name__ == "__main__":
32 | mercury_trajectory()
33 |
--------------------------------------------------------------------------------
/src/gros/metric/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/bytebat/gros/2e0e68a315249cb8707a344a0586ccd77a2f1b94/src/gros/metric/__init__.py
--------------------------------------------------------------------------------
/src/gros/metric/schwarzschild.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from astropy import units as u
3 | from scipy import integrate
4 |
5 | from gros.utils import const, log, datahandling as dh
6 |
7 |
8 | logger = log.init_logger(__name__)
9 |
10 |
11 | class SchwarzschildMetric:
12 | """
13 | Class for calculations with the Schwarschild metric,
14 | e.g. for solving the geodesic equation for an orbiting object.
15 |
16 | Note: See following lecture notes for a deeper look into the used formulas:
17 | https://web.stanford.edu/~oas/SI/SRGR/notes/SchwarzschildSolution.pdf
18 | https://arxiv.org/pdf/0904.4184.pdf
19 | http://www.physics.usu.edu/Wheeler/GenRel/Lectures/GRNotesDecSchwarzschildGeodesicsPost.pdf
20 | """
21 |
22 | @u.quantity_input(M=u.kg, time=u.s)
23 | def __init__(self, M, initial_vec_pos, initial_vec_v, time=0 * u.s):
24 | """
25 | Initializes a Schwarzschild metric with given mass,
26 | start time and spatial coordinates and velocities of a test particle.
27 |
28 | Arguments:
29 | M -- mass of gravitational center [kg]
30 | initial_vec_pos -- initial spherical coordinates [r[m], theta[rad], phi[rad]]
31 | sperical_velos -- initial spherical velocities [r', theta', phi']
32 | time -- initial time [s]
33 | """
34 | self.M = M.value
35 | self.rs = calc_schwarzschild_radius(M).value
36 | self.a = -self.rs
37 |
38 | logger.info("Event horizon for given mass is at r={}m.".format(self.rs))
39 |
40 | self.initial_vec_x_u = np.hstack(
41 | (
42 | time.value,
43 | initial_vec_pos,
44 | self._calc_initial_dt_dtau(initial_vec_pos, initial_vec_v),
45 | initial_vec_v,
46 | )
47 | )
48 |
49 | def schwarzschild_radius(self):
50 | """
51 | Returns the Schwarzschild radius for the initialized mass.
52 | """
53 | return self.rs
54 |
55 | def calculate_trajectory(self, proptime_end, proptime_start=0, step_size=1e-3):
56 | """
57 | Solve geodesic equations for provided time interval.
58 |
59 | Arguments:
60 | proptime_end -- end of interval [s]
61 | proptime_start -- start of inteval [s]
62 | step_size -- step size [s]
63 |
64 | Returns:
65 | (N,9) numpy ndarray with N = number of interations,
66 | single row is [tau, t, x, y, z]
67 | """
68 | tau = []
69 | y = []
70 |
71 | for cur_tau, cur_val in self.trajectory_generator(
72 | proptime_end=proptime_end,
73 | proptime_start=proptime_start,
74 | step_size=step_size,
75 | ):
76 | tau.append(cur_tau), y.append(cur_val)
77 |
78 | logger.debug(
79 | "tau={tau}s, t={t}s, r={r}m, dt/dtau={dtdtau}".format(
80 | tau=cur_tau, t=cur_val[0], r=cur_val[1], dtdtau=cur_val[4]
81 | )
82 | )
83 |
84 | np_tau = np.reshape(tau, (len(tau), 1))
85 | np_y = np.array(y)
86 |
87 | # only extract rows for [tau, t, x, y, z]
88 | return dh.SpaceTimeData(np.hstack((np_tau, np_y[:, [0, 1, 2, 3]])), self.rs)
89 |
90 | def trajectory_generator(self, proptime_end, proptime_start=0, step_size=1e-3):
91 | """
92 | Generator for solving the next time step of the geodesic's ODE system.
93 |
94 | Arguments:
95 | proptime_end -- end of interval [s]
96 | proptime_start -- start proper time [s]
97 | step_size -- step size [s]
98 |
99 | Yields:
100 | proper time,
101 | vector composed of the updated 4-vector and 4-velocity
102 | => [t, r, theta, phi, dt/dtau, dr/dtau, dtheta/dtau, dphi/dtau]
103 | """
104 | ode_solver = integrate.RK45(
105 | self._calc_next_velo_and_acc_vec,
106 | t0=proptime_start,
107 | y0=self.initial_vec_x_u,
108 | t_bound=1e100,
109 | rtol=0.25 * step_size,
110 | first_step=0.75 * step_size,
111 | max_step=5 * step_size,
112 | )
113 |
114 | while True:
115 | yield ode_solver.t, ode_solver.y
116 | ode_solver.step()
117 |
118 | current_radius = ode_solver.y[1]
119 | rs = -self.a
120 | rs_border = rs * 1.01
121 |
122 | if ode_solver.t >= proptime_end:
123 | break
124 |
125 | if current_radius <= rs_border:
126 | logger.warning(
127 | "Nearly reached event horizon at r={}m. Stopping here!".format(rs)
128 | )
129 | break
130 |
131 | def _calc_christoffel_symbols(self, r, theta):
132 | """
133 | Calculate the Christoffel symbols for the initialized mass
134 | at the coordinates r, theta.
135 |
136 | Arguments:
137 | r -- radius in m
138 | theta -- angle in radians (z-axis)
139 |
140 | Returns:
141 | Christoffel symbols as (4,4,4) numpy array
142 | """
143 | a = self.a
144 | a_div_r = a / r
145 | chrs = np.zeros(shape=(4, 4, 4), dtype=float)
146 |
147 | chrs[0, 0, 1] = chrs[0, 1, 0] = -0.5 * a / ((r ** 2) * (1 + a_div_r))
148 | chrs[1, 0, 0] = -0.5 * (1 + a_div_r) * (const.c.value ** 2) * a_div_r / r
149 | chrs[1, 1, 1] = -chrs[0, 0, 1]
150 | chrs[2, 1, 2] = chrs[2, 2, 1] = 1 / r
151 | chrs[1, 2, 2] = -1 * (r + a)
152 | chrs[3, 1, 3] = chrs[3, 3, 1] = chrs[2, 1, 2]
153 | chrs[1, 3, 3] = -1 * (r + a) * (np.sin(theta) ** 2)
154 | chrs[3, 2, 3] = chrs[3, 3, 2] = 1 / np.tan(theta)
155 | chrs[2, 3, 3] = -1 * np.sin(theta) * np.cos(theta)
156 | return chrs
157 |
158 | def _calc_initial_dt_dtau(self, vec_pos, vec_v):
159 | """
160 | Calculates the initial time velocity dt/dtau.
161 |
162 | Arguments:
163 | vec_pos -- spatial part of four-vector (r, theta, phi)
164 | vec_v -- spatial part of four-velocity (dr/dtau, dtheta/dtau, dphi/dtau)
165 |
166 | Returns:
167 | Initial time velocity dt/dtau
168 | """
169 | a, c = self.a, const.c.value
170 | c2 = c ** 2
171 |
172 | temp1 = (1 / (c2 * (1 + a / vec_pos[0]))) * (vec_v[0] ** 2)
173 | temp2 = ((vec_pos[0] ** 2) / c2) * (vec_v[1] ** 2)
174 | temp3 = (vec_pos[0] ** 2) / c2 * (np.sin(vec_pos[1]) ** 2) * (vec_v[2] ** 2)
175 | dt_dtau_squared = (1 + temp1 + temp2 + temp3) / (1 + a / vec_pos[0])
176 | return np.sqrt(dt_dtau_squared)
177 |
178 | def _calc_next_velo_and_acc_vec(self, proper_time, vec_x_u):
179 | """
180 | Calculates the acceleration components u' of a particle in the gravitational field
181 | based on the current position (four-vector x) and the corresponding four-velocity u.
182 |
183 | Arguments:
184 | vec_x_u -- combined four-vector x and four-velocity u,
185 | [x, u] = [t, r, theta, phi, dt/dtau, dr/dtau, dtheta/dtau, dphi/dtau]
186 | Returns:
187 | 8-component-vector [u, u'], composed of the derivatives of the input vector
188 | """
189 | derivs = np.zeros(shape=vec_x_u.shape, dtype=vec_x_u.dtype)
190 | chs = self._calc_christoffel_symbols(r=vec_x_u[1], theta=vec_x_u[2])
191 |
192 | derivs[:4] = vec_x_u[4:8]
193 | derivs[4] = -2 * chs[0, 0, 1] * vec_x_u[4] * vec_x_u[5]
194 | derivs[5] = -1 * (
195 | chs[1, 0, 0] * (vec_x_u[4] ** 2)
196 | + chs[1, 1, 1] * (vec_x_u[5] ** 2)
197 | + chs[1, 2, 2] * (vec_x_u[6] ** 2)
198 | + chs[1, 3, 3] * (vec_x_u[7] ** 2)
199 | )
200 | derivs[6] = -2 * chs[2, 2, 1] * vec_x_u[6] * vec_x_u[5] - chs[2, 3, 3] * (
201 | vec_x_u[7] ** 2
202 | )
203 | derivs[7] = -2 * (
204 | chs[3, 1, 3] * vec_x_u[5] * vec_x_u[7]
205 | + chs[3, 2, 3] * vec_x_u[6] * vec_x_u[7]
206 | )
207 |
208 | return derivs
209 |
210 |
211 | # TODO:
212 | # class Geodesic(SchwarzschildMetric):
213 | # """[summary]
214 |
215 | # Arguments:
216 | # SchwarzschildMetric {[type]} -- [description]
217 | # """
218 | # @u.quantity_input(M=u.kg, stop_time=u.s, start_time=u.s)
219 | # def __init__(self, M, initial_vec_pos, initial_vec_v, stop_time, start_time=0 * u.s):
220 | # super.__init__(M, initial_vec_pos, initial_vec_v, start_time)
221 |
222 |
223 | @u.quantity_input(M=u.kg)
224 | def calc_schwarzschild_radius(M=u.kg):
225 | """
226 | Calculates the Schwarzschild radius of a point mass M,
227 | which defines the event horizon of a Schwarzschild black hole.
228 |
229 | Arguments:
230 | M -- mass [kg]
231 |
232 | Returns:
233 | Schwarzschild radius [m]
234 | """
235 | return 2 * const.G * M / (const.c ** 2)
236 |
--------------------------------------------------------------------------------
/src/gros/tests/metric/test_schwarzschild.py:
--------------------------------------------------------------------------------
1 | from gros.metric import schwarzschild as ss
2 | from gros.utils import const
3 | from numpy import pi
4 | from numpy import sqrt
5 |
6 |
7 | def test_calc_schwarzschild_trajectory():
8 | r = 140000000000
9 |
10 | ss_metric = ss.SchwarzschildMetric(
11 | M=const.M_solar,
12 | initial_vec_pos=[r, pi/2, 0.3],
13 | initial_vec_v=[1000, 0.0, 0.0]
14 | )
15 |
16 | for val in ss_metric.trajectory_generator(proptime_end=10, step_size=1):
17 | print("tau = {}s, t = {}s, r = {}m, dt/dtau = {}".format(val[0], val[1][0], val[1][1], val[1][4]))
18 |
19 | dt_factor = 1/sqrt(1 - ss_metric.schwarzschild_radius() / r)
20 | print("theoretical dt/dtau = {}".format(dt_factor))
21 |
22 | assert(val[1][0] > val[0])
23 |
24 |
25 | test_calc_schwarzschild_trajectory()
26 |
--------------------------------------------------------------------------------
/src/gros/tests/utils/test_datahandling.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 | from gros.utils import datahandling as dh
4 |
5 |
6 | def test_general_data_handling():
7 | testdata = np.array([
8 | [0, 0, 1, 2, 3],
9 | [1, 1, 4, 5, 6]
10 | ])
11 | traj_data = dh.SpaceTimeData(testdata, 1000)
12 | assert(traj_data.size() == 2)
13 |
14 |
15 | test_general_data_handling()
16 |
--------------------------------------------------------------------------------
/src/gros/utils/__init__.py:
--------------------------------------------------------------------------------
1 | from .const import *
2 | from .datahandling import *
3 | from .log import *
4 | from .transforms import *
5 |
--------------------------------------------------------------------------------
/src/gros/utils/const.py:
--------------------------------------------------------------------------------
1 | from astropy import constants
2 |
3 |
4 | c = constants.c
5 | G = constants.G
6 | M_solar = constants.M_sun
7 | M_earth = constants.M_earth
8 |
--------------------------------------------------------------------------------
/src/gros/utils/datahandling.py:
--------------------------------------------------------------------------------
1 | import math
2 | import numpy as np
3 | import pandas as pd
4 | import plotly.graph_objects as go
5 |
6 | from gros.utils import log, transforms as tf
7 |
8 |
9 | logger = log.init_logger(__name__)
10 | DATA_COLUMNS = ["tau", "t", "x", "y", "z"]
11 |
12 |
13 | class SpaceTimeData:
14 | """
15 | Class for wrapping spacetime data points.
16 | """
17 |
18 | def __init__(
19 | self, trajectory, rs, conversion=tf.CoordinateConversion.spherical_to_cartesian
20 | ):
21 | """
22 | Initialize dataset of trajectory points.
23 | Arguments:
24 | rs -- Schwarzschild radius of center mass [m]
25 | """
26 | if not isinstance(trajectory, np.ndarray):
27 | raise ValueError("dataset must be numpy ndarray!")
28 | if trajectory.shape != (len(trajectory), 5):
29 | raise ValueError("dataset shape needs to be (N,5)!")
30 |
31 | self.rs = rs
32 | self.max_num_anim_frames = 100
33 |
34 | if trajectory.any():
35 | self.df = pd.DataFrame(trajectory, columns=DATA_COLUMNS)
36 | if conversion is tf.CoordinateConversion.spherical_to_cartesian:
37 | for traj_point in trajectory:
38 | (
39 | traj_point[2],
40 | traj_point[3],
41 | traj_point[4],
42 | ) = tf.spherical_to_cartesian(
43 | r=traj_point[2], theta=traj_point[3], phi=traj_point[4]
44 | )
45 |
46 | def size(self):
47 | """
48 | Returns number of datapoints.
49 | """
50 | return len(self.df.index)
51 |
52 | def plot(self, attractor_radius=0, animation_step_size=0, theme="dark"):
53 | """
54 | Plots the provided data points in a 3D scatter plot.
55 |
56 | Arguments:
57 | attractor_radius {int} -- Adds an additional sphere with given radius [m] to the plot (default: {0})
58 | animation_step_size {int} -- step size [s] > 0 will create an according animation (default: {0})
59 | theme {str} -- Plotting theme (default: {"dark"})
60 | """
61 | TRAJ_COLOR = "skyblue"
62 | SINGULARITY_COLOR = "darkviolet"
63 | BLACK_HOLE_COLOR = "darkviolet"
64 |
65 | data = self.df
66 |
67 | axis_min = min([data["x"].min(), data["y"].min(), data["z"].min()])
68 | axis_max = max([data["x"].max(), data["y"].max(), data["z"].max()])
69 |
70 | x = data["x"]
71 | y = data["y"]
72 | z = data["z"]
73 |
74 | # create marker for singularity
75 | singularity = go.Scatter3d(
76 | x=[0],
77 | y=[0],
78 | z=[0],
79 | mode="markers",
80 | marker=dict(size=1, color=SINGULARITY_COLOR),
81 | text="singularity",
82 | name="singularity",
83 | hoverinfo="text",
84 | )
85 |
86 | # attractor sphere
87 | attractor = self._create_sphere(attractor_radius, "attractor", "yellow", 0.1)
88 |
89 | # black hole sphere at r=rs
90 | black_hole = self._create_sphere(self.rs, "black hole", BLACK_HOLE_COLOR, 0.2)
91 |
92 | # trajectory
93 | traj_plot = go.Scatter3d(
94 | x=x,
95 | y=y,
96 | z=z,
97 | text="trajectory",
98 | name="trajectory",
99 | mode="lines",
100 | line=dict(width=2, color=TRAJ_COLOR),
101 | )
102 |
103 | plot_data = [
104 | go.Scatter3d(name="dummy"),
105 | singularity,
106 | black_hole,
107 | attractor,
108 | traj_plot,
109 | ]
110 |
111 | fig = go.Figure(
112 | data=plot_data,
113 | layout=go.Layout(
114 | scene=dict(
115 | xaxis=dict(range=[axis_min, axis_max],),
116 | yaxis=dict(range=[axis_min, axis_max],),
117 | zaxis=dict(range=[axis_min, axis_max],),
118 | aspectmode="cube",
119 | ),
120 | title_text="Particle orbit in gravitational field",
121 | hovermode="closest",
122 | updatemenus=[],
123 | ),
124 | )
125 |
126 | self._add_aninmation_frames(fig, animation_step_size)
127 |
128 | if theme == "dark":
129 | fig.update_layout(template="plotly_dark")
130 | fig.show()
131 |
132 | def _create_sphere(self, radius, name, color, opacity):
133 | sph_theta = np.linspace(0, 2 * np.pi)
134 | sph_phi = np.linspace(0, 2 * np.pi)
135 | sph_theta, sph_phi = np.meshgrid(sph_phi, sph_theta)
136 |
137 | sph_x, sph_y, sph_z = tf.spherical_to_cartesian(radius, sph_theta, sph_phi)
138 |
139 | sphere = go.Mesh3d(
140 | x=sph_x.flatten(),
141 | y=sph_y.flatten(),
142 | z=sph_z.flatten(),
143 | alphahull=0,
144 | opacity=opacity,
145 | color=color,
146 | text=name,
147 | name=name,
148 | )
149 |
150 | return sphere
151 |
152 | def _add_aninmation_frames(self, fig, anim_step_size):
153 | """Adds animation frames with given (index) step size to the figure.
154 |
155 | Arguments:
156 | fig -- plotly figure
157 | anim_step_size -- animation step size [s]
158 | """
159 | if anim_step_size <= 0:
160 | return
161 |
162 | # buttons
163 | updatemenus_dict = {
164 | # "bgcolor": "#333333",
165 | "font": {"color": "#000000"},
166 | "type": "buttons",
167 | "buttons": [
168 | {
169 | "label": "Play",
170 | "method": "animate",
171 | "args": [
172 | None,
173 | {
174 | "frame": {"duration": 500, "redraw": True},
175 | "fromcurrent": True,
176 | "transition": {
177 | "duration": 300,
178 | "easing": "quadratic-in-out",
179 | },
180 | },
181 | ],
182 | },
183 | {
184 | "label": "Pause",
185 | "method": "animate",
186 | "args": [
187 | [None],
188 | {
189 | "frame": {"duration": 0, "redraw": False},
190 | "mode": "immediate",
191 | "transition": {"duration": 0},
192 | },
193 | ],
194 | },
195 | ],
196 | "showactive": False,
197 | }
198 |
199 | # animation frame traits
200 | frame_step_size = math.ceil(anim_step_size / self.df["tau"].max() * self.size())
201 |
202 | if frame_step_size > self.max_num_anim_frames:
203 | frame_step_size = self.max_num_anim_frames
204 | logger.warning(
205 | "'animation_step_size' is too large. \
206 | Maximium number of animation frames will be limited to {}.".format(
207 | self.max_num_anim_frames
208 | )
209 | )
210 |
211 | # sliders
212 | sliders_dict = {
213 | "active": 0,
214 | "yanchor": "top",
215 | "xanchor": "left",
216 | "currentvalue": {
217 | "font": {"size": 12},
218 | "prefix": "tau, t:
3 | 
60 | 
68 |