Merge branch 'main' into update_readme

docs/changes/46.feature.rst

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

docs/changes/48.feature.rst

@@ -0,0 +1 @@
+- Added optional auto scaling for batchsize in vis_loop

environment.yml

@@ -16,3 +16,5 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-runner
+  - pip:
+    - toma

pyproject.toml

@@ -28,25 +28,26 @@ classifiers = [
 requires-python = ">=3.10"
 
 dependencies = [
-  "numpy",
+  "astroplan",
   "astropy<=6.1.0",
-  "torch",
-  "matplotlib",
+  "click",
+  "h5py",
   "ipython",
-  "scipy",
+  "jupyter",
+  "matplotlib",
+  "natsort",
+  "numexpr",
+  "numpy",
   "pandas",
-  "toml",
+  "pre-commit",
   "pytest",
   "pytest-cov",
-  "jupyter",
-  "astroplan",
+  "scipy",
+  "toma",
+  "toml",
+  "torch",
   "torch",
   "tqdm",
-  "numexpr",
-  "click",
-  "h5py",
-  "natsort",
-  "pre-commit",
 ]
 
 [project.scripts]

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)

pyvisgen/simulation/visibility.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass, fields
 
 import torch
-from tqdm import tqdm
+import toma
+from tqdm.autonotebook import tqdm
 
 import pyvisgen.simulation.scan as scan
 
@@ -44,12 +45,18 @@ def vis_loop(
     num_threads=10,
     noisy=True,
     mode="full",
-    batch_size=100,
+    batch_size="auto",
     show_progress=False,
 ):
     torch.set_num_threads(num_threads)
     torch._dynamo.config.suppress_errors = True
 
+    if not (
+        isinstance(batch_size, int)
+        or (isinstance(batch_size, str) and batch_size == "auto")
+    ):
+        raise ValueError("Expected batch_size to be 'auto' or of type int")
+
     SI = torch.flip(SI, dims=[1])
 
     # define unpolarized sky distribution
@@ -104,10 +111,78 @@ def vis_loop(
     else:
         raise ValueError("Unsupported mode!")
 
-    batches = torch.arange(bas[:].shape[1]).split(batch_size)
+    if batch_size == "auto":
+        batch_size = bas[:].shape[1]
+
+    visibilities = toma.explicit.batch(
+        _batch_loop,
+        batch_size,
+        visibilities,
+        vis_num,
+        obs,
+        B,
+        bas,
+        lm,
+        rd,
+        noisy,
+        show_progress,
+    )
+
+    return visibilities
 
-    if show_progress:
-        batches = tqdm(batches)
+
+def _batch_loop(
+    batch_size: int,
+    visibilities,
+    vis_num: int,
+    obs,
+    B: torch.tensor,
+    bas,
+    lm: torch.tensor,
+    rd: torch.tensor,
+    noisy: bool | float,
+    show_progress: bool,
+):
+    """Main simulation loop of pyvisgen. Computes visibilities
+    batchwise.
+
+    Parameters
+    ----------
+    batch_size : int
+        Batch size for loop over Baselines dataclass object.
+    visibilities : Visibilities
+        Visibilities dataclass object.
+    vis_num : int
+        Number of visibilities.
+    obs : Observation
+        Observation class object.
+    B : torch.tensor
+        Stokes matrix containing stokes visibilities.
+    bas : Baselines
+        Baselines dataclass object.
+    lm : torch.tensor
+        lm grid.
+    rd : torch.tensor
+        rd grid.
+    noisy : float or bool
+        Simulate noise as SEFD with given value. If set to False,
+        no noise is simulated.
+    show_progress :
+        If True, show a progress bar tracking the loop.
+
+    Returns
+    -------
+    visibilities : Visibilities
+        Visibilities dataclass object.
+    """
+    batches = torch.arange(bas[:].shape[1]).split(batch_size)
+    batches = tqdm(
+        batches,
+        position=0,
+        disable=not show_progress,
+        desc="Computing visibilities",
+        postfix=f"Batch size: {batch_size}",
+    )
 
     for p in batches:
         bas_p = bas[:][:, p]
@@ -155,6 +230,7 @@ def vis_loop(
 
         visibilities.add(vis)
         del int_values
+
     return visibilities
 
 

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())