diff --git a/docs/changes/46.feature.rst b/docs/changes/46.feature.rst new file mode 100644 index 0000000..3739cc3 --- /dev/null +++ b/docs/changes/46.feature.rst @@ -0,0 +1 @@ +- ``pyvisgen.layouts.get_array_layout`` now also accepts custom layouts stored in a ``pd.DataFrame`` diff --git a/pyvisgen/layouts/layouts.py b/pyvisgen/layouts/layouts.py index a88b68d..03ebdc4 100644 --- a/pyvisgen/layouts/layouts.py +++ b/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) diff --git a/tests/data/test_layout.txt b/tests/data/test_layout.txt new file mode 100644 index 0000000..f37e305 --- /dev/null +++ b/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 diff --git a/tests/test_layouts.py b/tests/test_layouts.py index 5bee463..20ac5b0 100644 --- a/tests/test_layouts.py +++ b/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())