Add support for more varieties of the telescope model (.tm)

scripts/misc/add_beaminfo.py

@@ -10,12 +10,15 @@ from numpy.linalg import norm
 import numpy as np
 import logging
 import sys
-
+import os
+from glob import glob
 
 def read_telescope_center(oskar_telescope_dir: str):
     """Read telescope center from OSKAR metadata and convert to ITRF"""
+
+    # separator is either comma or whitespace, plus any extra whitespace
     oskar_telescope_center = pd.read_csv(f"{oskar_telescope_dir}/position.txt",
-                                         header=None, sep=" ")
+                                         header=None, sep="[,\s]\s*")
 
     if len(oskar_telescope_center.columns) == 2:
         oskar_telescope_center["Height"] = [0.]
@@ -24,8 +27,22 @@ def read_telescope_center(oskar_telescope_dir: str):
     telescope_center_itrf = geo.xyz_from_geographic(np.deg2rad(oskar_telescope_center["Lon"][0]),
                                                     np.deg2rad(oskar_telescope_center["Lat"][0]),
                                                     oskar_telescope_center["Height"][0])
-    return telescope_center_itrf
 
+    normal_vector_ellipsoid = geo.normal_vector_ellipsoid(np.deg2rad(oskar_telescope_center["Lon"][0]),
+                                                          np.deg2rad(oskar_telescope_center["Lat"][0]))
+
+    local_to_itrf_projection_matrix = geo.projection_matrix(telescope_center_itrf,
+                                                            normal_vector_ellipsoid)
+    return telescope_center_itrf, local_to_itrf_projection_matrix
+
+def read_element_locations(oskar_layout_file: str):
+    # separator is either comma or whitespace, plus any extra whitespace
+    element_locations = pd.read_csv(oskar_layout_file,
+                                    header=None, sep="[,\s]\s*")
+    if len(element_locations.columns) == 2:
+        element_locations["Z"] = np.zeros_like(element_locations[0])
+
+    return element_locations
 
 def fix_antenna(oskar_ms_name: str, telescope_center_itrf: np.array):
     """Make POSITION in ::ANTENNA subtable absolute"""
@@ -93,29 +110,26 @@ def add_phased_array_table(oskar_ms_name: str):
 
 def fill_phased_array(oskar_ms_name: str, oskar_telescope_dir: str):
     """Fill the ::PHASED_ARRAY subtable with info from the OSKAR directory"""
-    element_locations = pd.read_csv(f"{oskar_telescope_dir}/station/layout.txt",
-                                    header=None, sep=" ")
-    if len(element_locations.columns) == 2:
-        element_locations["Z"] = np.zeros_like(element_locations[0])
-
-    oskar_telescope_center = pd.read_csv(f"{oskar_telescope_dir}/position.txt",
-                                         header=None, sep=" ")
-    oskar_telescope_center.columns = ["Lon", "Lat"]
-    telescope_center_itrf = read_telescope_center(oskar_telescope_dir)
 
-    normal_vector_ellipsoid = geo.normal_vector_ellipsoid(np.deg2rad(oskar_telescope_center["Lon"][0]),
-                                                          np.deg2rad(oskar_telescope_center["Lat"][0]))
+    telescope_center_itrf, local_to_itrf_projection_matrix = read_telescope_center(oskar_telescope_dir)
 
-    local_to_itrf_projection_matrix = geo.projection_matrix(telescope_center_itrf,
-                                                            normal_vector_ellipsoid)
+    element_layout_file = f"{oskar_telescope_dir}/station/layout.txt"
+    if os.path.exists(element_layout_file):
+        element_locations = read_element_locations(element_layout_file)
 
-    element_locations_itrf = local_to_itrf_projection_matrix @ element_locations.to_numpy().T
     nr_pol = 2
-    all_unflagged = np.zeros((nr_pol, element_locations_itrf.shape[1]))
 
     phasedarraytable = pt.table(f"{oskar_ms_name}::PHASED_ARRAY", readonly=False, ack=False)
  
-    for stationnr in range(len(phasedarraytable)):
+    oskar_station_dirs = sorted(glob(f"{oskar_telescope_dir}/station*"))
+
+    for stationnr, oskar_station_dir in enumerate(oskar_station_dirs):
+        element_layout_file = f"{oskar_station_dir}/layout.txt"
+        print(element_layout_file)
+        if os.path.exists(element_layout_file):
+            element_locations = read_element_locations(element_layout_file)
+        element_locations_itrf = local_to_itrf_projection_matrix @ element_locations.to_numpy().T
+        all_unflagged = np.zeros((nr_pol, element_locations_itrf.shape[1]))
         phasedarraytable.putcell("COORDINATE_AXES", stationnr, local_to_itrf_projection_matrix.T)
         phasedarraytable.putcell("ELEMENT_OFFSET", stationnr, element_locations_itrf)
         phasedarraytable.putcell("ELEMENT_FLAG", stationnr, all_unflagged)
@@ -127,7 +141,7 @@ def fill_phased_array(oskar_ms_name: str, oskar_telescope_dir: str):
 
 def main(oskar_ms_name: str, oskar_telescope_dir: str):
     """Add beam info to OSKAR generated MeasurementSet"""
-    telescope_center_itrf = read_telescope_center(oskar_telescope_dir)
+    telescope_center_itrf,_ = read_telescope_center(oskar_telescope_dir)
     fix_antenna(oskar_ms_name, telescope_center_itrf)
     add_array_center(oskar_ms_name, telescope_center_itrf)
     add_phased_array_table(oskar_ms_name)