Allow custom layouts via DataFrames in pyvisgen.layouts.get_array_layout

docs/changes/46.feature.rst

@@ -0,0 +1 @@
+- ``pyvisgen.layouts.get_array_layout`` now also accepts custom layouts stored in a ``pd.DataFrame``

pyvisgen/layouts/layouts.py

@@ -24,34 +24,46 @@ def __getitem__(self, i):
         return Stations(*[getattr(self, f.name)[i] for f in fields(self)])
 
 
-def get_array_layout(array_name, writer=False):
+def get_array_layout(array_layout: str | Path | pd.DataFrame, writer: bool = False):
     """Reads telescope layout txt file and converts it into a dataclass.
+    Also allows a DataFrame to be passed that is then converted into a dataclass
+    object.
     Available arrays:
     - EHT
 
     Parameters
     ----------
-    array_name : str
-        Name of telescope array
+    array_layout : str or pathlib.Path or pd.DataFrame
+        Name of telescope array or pd.DataFrame containing
+        the array layout.
+    writer : bool, optional
+        If ``True``, return ``array`` DataFrame instead of
+        ``Stations`` dataclass object.
 
     Returns
     -------
     dataclass objects
-        Station infos combinde in dataclass
+        Station infos combined in dataclass
     """
-    f = array_name + ".txt"
-    array = pd.read_csv(file_dir / f, sep=r"\s+")
-    if array_name == "vla":
-        loc = EarthLocation.of_site("VLA")
-        array["X"] += loc.value[0]
-        array["Y"] += loc.value[1]
-        array["Z"] += loc.value[2]
-
-    if array_name == "test_layout":
-        loc = EarthLocation.of_address("dortmund")
-        array["X"] += loc.value[0]
-        array["Y"] += loc.value[1]
-        array["Z"] += loc.value[2]
+    if isinstance(array_layout, str):
+        f = array_layout + ".txt"
+        array = pd.read_csv(file_dir / f, sep=r"\s+")
+
+        if array_layout == "vla":
+            # Change relative positions to absolute positions
+            # for the VLA layout
+            loc = EarthLocation.of_site("VLA")
+            array["X"] += loc.value[0]
+            array["Y"] += loc.value[1]
+            array["Z"] += loc.value[2]
+
+    elif isinstance(array_layout, pd.DataFrame):
+        array = array_layout
+    else:
+        raise TypeError(
+            "Expected array_layout to be of type str, "
+            "pathlib.Path, or pandas.DataFrame!"
+        )
 
     # drop name col and convert to tensor
     tensor = torch.from_numpy(array.iloc[:, 1:].values)

tests/data/test_layout.txt

@@ -0,0 +1,11 @@
+station_name X Y Z dish_dia el_low el_high SEFD altitude
+t000 -4000 2000 0 25 15.0 85.0 110.0 1000.0
+t001 8000 -2000 0 25 15.0 85.0 110.0 1000.0
+t002 1000 1000 0 25 15.0 85.0 110.0 1000.0
+t003 -4000 6000 0 25 15.0 85.0 110.0 1000.0
+t004 -3000 -3000 0 25 15.0 85.0 110.0 1000.0
+t005 6000 -5000 0 25 15.0 85.0 110.0 1000.0
+t006 8000 2000 0 25 15.0 85.0 110.0 1000.0
+t007 2000 -4000 0 25 15.0 85.0 110.0 1000.0
+t008 -6000 3000 0 25 15.0 85.0 110.0 1000.0
+t009 -2000 8000 0 25 15.0 85.0 110.0 1000.0

tests/test_layouts.py

@@ -1,26 +1,87 @@
 import torch
+from numpy.testing import assert_array_equal, assert_raises
 
+from pyvisgen.layouts.layouts import get_array_layout
 
-def test_get_array_layout():
-    from pyvisgen.layouts.layouts import get_array_layout
 
-    layout = get_array_layout("eht")
+class TestLayouts:
+    def setup_class(self):
+        from pathlib import Path
 
-    assert len(layout.st_num) == 8
-    assert torch.is_tensor(layout[0].x)
-    assert torch.is_tensor(layout[0].y)
-    assert torch.is_tensor(layout[0].z)
-    assert torch.is_tensor(layout[0].diam)
-    assert torch.is_tensor(layout[0].el_low)
-    assert torch.is_tensor(layout[0].el_high)
-    assert torch.is_tensor(layout[0].sefd)
-    assert torch.is_tensor(layout[0].altitude)
+        import pandas as pd
+        from astropy.coordinates import EarthLocation
 
-    layout = get_array_layout("vlba")
+        # Declare a reduced EHT test layout. This should suffice for
+        # testing purposes
+        self.eht_test_layout = pd.DataFrame(
+            {
+                "st_num": [0, 1, 2],
+                "x": [2225037.1851, -1828796.2, -768713.9637],
+                "y": [-5441199.162, -5054406.8, -5988541.7982],
+                "z": [-2479303.4629, 3427865.2, 2063275.9472],
+                "diam": [84.700000, 10.000000, 50.000000],
+                "el_low": [15.000000, 15.000000, 15.000000],
+                "el_high": [85.000000, 85.000000, 85.000000],
+                "sefd": [110.0, 11900.0, 560.0],
+                "altitude": [5030.0, 3185.0, 4640.0],
+            }
+        )
 
-    assert len(layout.st_num) == 10
+        # Load test layout from file
+        self.test_layout = pd.read_csv(
+            Path(__file__).parent.resolve() / "data/test_layout.txt", sep=r"\s+"
+        )
 
-    layout = get_array_layout("vla")
+        # Place test layout in Dortmund
+        loc = EarthLocation.of_address("dortmund")
+        self.test_layout["X"] += loc.value[0]
+        self.test_layout["Y"] += loc.value[1]
+        self.test_layout["Z"] += loc.value[2]
 
-    assert len(layout.st_num) == 27
-    assert layout[:3].st_num.shape == torch.Size([3])
+    def test_get_array_layout(self):
+        layout = get_array_layout("eht")
+
+        assert len(layout.st_num) == 8
+
+        assert_array_equal(layout[:3].x, self.eht_test_layout.x)
+        assert_array_equal(layout[:3].y, self.eht_test_layout.y)
+        assert_array_equal(layout[:3].z, self.eht_test_layout.z)
+        assert_array_equal(layout[:3].diam, self.eht_test_layout.diam)
+        assert_array_equal(layout[:3].el_low, self.eht_test_layout.el_low)
+        assert_array_equal(layout[:3].el_high, self.eht_test_layout.el_high)
+        assert_array_equal(layout[:3].sefd, self.eht_test_layout.sefd)
+        assert_array_equal(layout[:3].altitude, self.eht_test_layout.altitude)
+
+        layout = get_array_layout("vlba")
+
+        assert len(layout.st_num) == 10
+
+        layout = get_array_layout("vla")
+
+        assert len(layout.st_num) == 27
+        assert layout[:3].st_num.shape == torch.Size([3])
+
+    def test_get_array_layout_dataframe(self):
+        layout = get_array_layout(self.test_layout)
+
+        assert_array_equal(layout.x, self.test_layout.X)
+        assert_array_equal(layout.y, self.test_layout.Y)
+        assert_array_equal(layout.z, self.test_layout.Z)
+        assert_array_equal(layout.diam, self.test_layout.dish_dia)
+        assert_array_equal(layout.el_low, self.test_layout.el_low)
+        assert_array_equal(layout.el_high, self.test_layout.el_high)
+        assert_array_equal(layout.sefd, self.test_layout.SEFD)
+        assert_array_equal(layout.altitude, self.test_layout.altitude)
+
+    def test_get_array_layout_raise(self):
+        # Converting the test DataFrame to a dict should raise
+        # a TypeError, since only str, pathlib.Path or pd.DataFrames
+        # are allowed
+        assert_raises(TypeError, get_array_layout, self.test_layout.to_dict)
+
+    def test_get_array_names(self):
+        from pyvisgen.layouts.layouts import get_array_names
+
+        test = ["vla", "vlba", "eht", "alma"]
+
+        assert set(test).issubset(get_array_names())