From 3acdd9b8aa4bec874a9d0966002e3f5cfcf57f9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20Fr=C3=B6se?= <stefan.froese@udo.edu>
Date: Wed, 25 Oct 2023 17:08:10 +0200
Subject: [PATCH] subarray takes EarthLocations as tel pos

---
 ctapipe/instrument/subarray.py            | 22 +++++++++----
 ctapipe/instrument/tests/test_subarray.py | 38 +++++++++++++++++++++++
 2 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/ctapipe/instrument/subarray.py b/ctapipe/instrument/subarray.py
index 3382245a9a7..f630f0bd95c 100644
--- a/ctapipe/instrument/subarray.py
+++ b/ctapipe/instrument/subarray.py
@@ -81,8 +81,8 @@ def __init__(
         ----------
         name : str
             name of this subarray
-        tel_positions : Dict[int, np.ndarray]
-            dict of x,y,z telescope positions on the ground by tel_id. These are
+        tel_positions : Dict[int, Union[np.ndarray, EarthLocation]]
+            dict of x,y,z telescope positions on the ground by tel_id or EarthLocation. These are
             converted internally to a coordinate in the `~ctapipe.coordinates.GroundFrame`
         tel_descriptions : Dict[TelescopeDescription]
             dict of TelescopeDescriptions by tel_id
@@ -91,7 +91,9 @@ def __init__(
             coordinate system used for `tel_positions`.
         """
         self.name = name
-        self.positions = tel_positions or dict()
+        self.positions: Dict[int, Union[np.ndarray, EarthLocation]] = (
+            tel_positions or dict()
+        )
         self.tels: Dict[int, TelescopeDescription] = tel_descriptions or dict()
         self.reference_location = reference_location
 
@@ -157,9 +159,17 @@ def info(self, printer=print):
     def tel_coords(self):
         """Telescope positions in `~ctapipe.coordinates.GroundFrame`"""
 
-        pos_x = [p[0].to_value(u.m) for p in self.positions.values()]
-        pos_y = [p[1].to_value(u.m) for p in self.positions.values()]
-        pos_z = [p[2].to_value(u.m) for p in self.positions.values()]
+        positions = self.positions.values()
+        positions = [
+            GroundFrame.from_earth_location(p, self.reference_location).cartesian.xyz
+            if isinstance(p, EarthLocation)
+            else p
+            for p in positions
+        ]
+
+        pos_x = [p[0].to_value(u.m) for p in positions]
+        pos_y = [p[1].to_value(u.m) for p in positions]
+        pos_z = [p[2].to_value(u.m) for p in positions]
 
         frame = GroundFrame(reference_location=self.reference_location)
 
diff --git a/ctapipe/instrument/tests/test_subarray.py b/ctapipe/instrument/tests/test_subarray.py
index 488cf6fc4fe..cf422d088a3 100644
--- a/ctapipe/instrument/tests/test_subarray.py
+++ b/ctapipe/instrument/tests/test_subarray.py
@@ -287,3 +287,41 @@ def test_subarrays(subarray_prod5_paranal: SubarrayDescription):
     assert subarray.name == "NewArray"
     assert isinstance(subarray.reference_location, EarthLocation)
     assert subarray.reference_location == subarray_prod5_paranal.reference_location
+
+
+def test_tel_pos_from_EarthLocation(prod5_mst_nectarcam):
+    rng = np.random.default_rng(0)
+
+    pos = {}
+    tel = {}
+
+    for tel_id in range(1, 11):
+        tel[tel_id] = prod5_mst_nectarcam
+        pos[tel_id] = rng.uniform(-100, 100, size=3) * u.m
+
+    for tel_id in range(12, 23):
+        tel[tel_id] = prod5_mst_nectarcam
+        rnd_lon = rng.uniform(
+            LOCATION.lon.to_value("deg") - 1e-3, LOCATION.lon.to_value("deg") + 1e-3
+        )
+        rnd_lat = rng.uniform(
+            LOCATION.lat.to_value("deg") - 1e-3, LOCATION.lat.to_value("deg") + 1e-3
+        )
+        rnd_height = rng.uniform(
+            LOCATION.height.to_value("m") - 1e2, LOCATION.height.to_value("m") + 1e2
+        )
+        pos[tel_id] = EarthLocation(
+            lon=rnd_lon * u.deg, lat=rnd_lat * u.deg, height=rnd_height
+        )
+
+    subarray = SubarrayDescription(
+        "test array",
+        tel_positions=pos,
+        tel_descriptions=tel,
+        reference_location=LOCATION,
+    )
+
+    x, y, z = subarray.tel_coords.cartesian.xyz
+    assert all(x.value > -100) and all(x.value < 100)
+    assert all(y.value > -100) and all(y.value < 100)
+    assert all(z.value > -100) and all(z.value < 100)