Add ability to pass custom metrics to geodesic code.

einsteinpy · Nov 14, 2022 · b92e66f · b92e66f
1 parent 705b351
commit b92e66f
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 6 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -12,6 +12,7 @@ CHANGES
 CONTRIBUTORS
 - JeS24
 - mshumayl
+- a3ahmad
 
 CHANGES
 - [#607]: Fixed #586 #598 #599 #600 #601, Fixed a bug in tensor_lambdify
@@ -21,6 +22,7 @@ CHANGES
 - [#612]: Partially addressed #610, Added support for setting plot aspect ratio and title via API
 - [#619]: General codebase maintenance
 - [#620]: Fixed #540, #541, Partially addressed #614, Updated documentation
+- [#625]: Add ability to compute geodesics for arbitrary metrics
 
 0.4.0
 -----

diff --git a/src/einsteinpy/geodesic/geodesic.py b/src/einsteinpy/geodesic/geodesic.py
@@ -99,13 +99,18 @@ def __init__(
         }
 
         if metric not in _METRICS:
-            raise NotImplementedError(
-                f"'{metric}' is unsupported. Currently, these metrics are supported:\
-                \n1. Schwarzschild\n2. Kerr\n3. KerrNewman"
-            )
+            if not callable(metric):
+                raise NotImplementedError(
+                    f"'{metric}' is unsupported. Currently, these metrics are supported:\
+                    \n1. Schwarzschild\n2. Kerr\n3. KerrNewman\n4. Any user-supplied callable metric function"
+                )
+
+            self.metric_name = metric.__name__
+            self.metric = metric
+        else:
+            self.metric_name = metric
+            self.metric = _METRICS[metric]
 
-        self.metric_name = metric
-        self.metric = _METRICS[metric]
         self.metric_params = metric_params
         if metric == "Schwarzschild":
             self.metric_params = (0.0,)

diff --git a/tests/test_geodesic/test_geodesic.py b/tests/test_geodesic/test_geodesic.py
@@ -5,6 +5,7 @@
 from numpy.testing import assert_allclose
 
 from einsteinpy.geodesic import Geodesic, Nulllike, Timelike
+from einsteinpy.geodesic.utils import _kerr
 
 
 @pytest.fixture()
@@ -169,3 +170,27 @@ def test_kerr0_eq_kn00():
 
     assert_allclose(k.trajectory[0], kn.trajectory[0], atol=1e-6, rtol=1e-6)
     assert_allclose(k.trajectory[1], kn.trajectory[1], atol=1e-6, rtol=1e-6)
+
+
+def test_custom_callable_metric():
+    metric_params = (0.,)
+    q0 = [4., np.pi / 2, 0.]
+    p0 = [0., 0., 0.]
+
+    try:
+        c = Timelike(
+            metric=_kerr,
+            metric_params=metric_params,
+            position=q0,
+            momentum=p0,
+            steps=50,
+            delta=0.5,
+            return_cartesian=True,
+            suppress_warnings=True,
+        )
+        assert c.metric_name == _kerr.__name__
+        assert c.metric == _kerr
+
+    except:
+        assert False
+