Add dataclass attributes docstring, add Polarisation class to __init__

Author: aknierim

pyvisgen/simulation/__init__.py

@@ -11,12 +11,13 @@
 )
 from .observation import Baselines, Observation, ValidBaselineSubset
 from .scan import angularDistance, calc_beam, calc_fourier, integrate, jinc, rime
-from .visibility import Visibilities, generate_noise, vis_loop
+from .visibility import Polarisation, Visibilities, generate_noise, vis_loop
 
 __all__ = [
     "Array",
     "Baselines",
     "Observation",
+    "Polarisation",
     "ValidBaselineSubset",
     "Visibilities",
     "angularDistance",

pyvisgen/simulation/observation.py

@@ -6,14 +6,13 @@
 import astropy.units as un
 import torch
 from astropy.constants import c
-from astropy.coordinates import AltAz, Angle, EarthLocation, SkyCoord, Longitude
+from astropy.coordinates import AltAz, Angle, EarthLocation, Longitude, SkyCoord
 from astropy.time import Time
 from tqdm.autonotebook import tqdm
 
 from pyvisgen.layouts import layouts
 from pyvisgen.simulation.array import Array
 
-
 DEFAULT_POL_KWARGS = {
     "delta": 0,
     "amp_ratio": 0.5,
@@ -30,6 +29,21 @@
 
 @dataclass
 class Baselines:
+    """Baselines dataclass.
+
+    Attributes
+    ----------
+    st1 : :func:`~torch.tensor`
+    st2 : :func:`~torch.tensor`
+    u : :func:`~torch.tensor`
+    v : :func:`~torch.tensor`
+    w : :func:`~torch.tensor`
+    valid : :func:`~torch.tensor`
+    time : :func:`~torch.tensor`
+    q1 : :func:`~torch.tensor`
+    q2 : :func:`~torch.tensor`
+    """
+
     st1: torch.tensor
     st2: torch.tensor
     u: torch.tensor
@@ -112,6 +126,30 @@ def get_valid_subset(self, num_baselines, device):
 
 @dataclass()
 class ValidBaselineSubset:
+    """Valid baselines subset dataclass.
+
+    Attributes
+    ----------
+    u_start : :func:`~torch.tensor`
+    u_stop : :func:`~torch.tensor`
+    u_valid : :func:`~torch.tensor`
+    v_start : :func:`~torch.tensor`
+    v_stop : :func:`~torch.tensor`
+    v_valid : :func:`~torch.tensor`
+    w_start : :func:`~torch.tensor`
+    w_stop : :func:`~torch.tensor`
+    w_valid : :func:`~torch.tensor`
+    baseline_nums : :func:`~torch.tensor`
+    date : :func:`~torch.tensor`
+    q1_start : :func:`~torch.tensor`
+    q1_stop : :func:`~torch.tensor`
+    q1_valid : :func:`~torch.tensor`
+    q2_start : :func:`~torch.tensor`
+    q2_stop : :func:`~torch.tensor`
+    q2_valid : :func:`~torch.tensor`
+
+    """
+
     u_start: torch.tensor
     u_stop: torch.tensor
     u_valid: torch.tensor
@@ -207,6 +245,74 @@ def _lexsort(self, a, dim=-1):
 
 
 class Observation:
+    """Main observation simulation class.
+    The :class:`~pyvisgen.simulation.Observation` class
+    simulates the baselines and time steps during the
+    observation.
+
+    Parameters
+    ----------
+    src_ra : float
+        Source right ascension coordinate.
+    src_dec : float
+        Source declination coordinate.
+    start_time : datetime
+        Observation start time.
+    scan_duration : int
+        Scan duration.
+    num_scans : int
+        Number of scans.
+    scan_separation : int
+        Scan separation.
+    integration_time : int
+        Integration time.
+    ref_frequency : float
+        Reference frequency.
+    frequency_offsets : list
+        Frequency offsets.
+    bandwidths : list
+        Frequency bandwidth.
+    fov : float
+        Field of view.
+    image_size : int
+        Image size of the sky distribution.
+    array_layout : str
+        Name of an existing array layout. See :mod:`~pyvisgen.layouts`.
+    corrupted : bool
+        If ``True``, apply corruption during the vis loop.
+    device : str
+        Torch device to select for computation.
+    dense : bool, optional
+        If ``True``, apply dense baseline calculation of a perfect
+        interferometer. Default: ``False``
+    sensitivity_cut : float, optional
+        Sensitivity threshold, where only pixels above the value
+        are kept. Default: ``1e-6``
+    polarisation : str, optional
+        Choose between ``'linear'`` or ``'circular'`` or ``None`` to
+        simulate different types of polarisations or disable
+        the simulation of polarisation. Default: ``None``
+    pol_kwargs : dict, optional
+        Additional keyword arguments for the simulation
+        of polarisation. Default:
+        ``{'delta': 0,'amp_ratio': 0.5,'random_state': 42}``
+    field_kwargs : dict, optional
+        Additional keyword arguments for the random polarisation
+        field that is applied when simulating polarisation.
+        Default:
+        ``{'order': [1, 1],'scale': [0, 1],'threshold': None,'random_state': 42}``
+    show_progress : bool, optional
+        If ``True``, show a progress bar during the iteration over the
+        scans. Default: ``False``
+
+    Notes
+    -----
+    See :class:`~pyvisgen.simulation.Polarisation` and
+    :class:`~pyvisgen.simulation.Polarisation.rand_polarisation_field`
+    for more information on the keyword arguments in ``pol_kwargs``
+    and ``field_kwargs``, respectively.
+    """
+
     def __init__(
         self,
         src_ra: float,
@@ -260,21 +366,21 @@ def __init__(
         image_size : int
             Image size of the sky distribution.
         array_layout : str
-            Name of an existing array layout. See `~pyvisgen.layouts`.
+            Name of an existing array layout. See :mod:`~pyvisgen.layouts`.
         corrupted : bool
-            If `True`, apply corruption during the vis loop.
+            If ``True``, apply corruption during the vis loop.
         device : str
             Torch device to select for computation.
         dense : bool, optional
-            If `True`, apply dense baseline calculation of a perfect
-            interferometer. Default: `False`
+            If ``True``, apply dense baseline calculation of a perfect
+            interferometer. Default: ``False``
         sensitivity_cut : float, optional
             Sensitivity threshold, where only pixels above the value
-            are kept. Default: 1e-6
+            are kept. Default: ``1e-6``
         polarisation : str, optional
-            Choose between `'linear'` or `'circular'` or `None` to
+            Choose between ``'linear'`` or ``'circular'`` or ``None`` to
             simulate different types of polarisations or disable
-            the simulation of polarisation. Default: `None`
+            the simulation of polarisation. Default: ``None``
         pol_kwargs : dict, optional
             Additional keyword arguments for the simulation
             of polarisation. Default: `{
@@ -292,15 +398,15 @@ def __init__(
                 "random_state": 42
             }`
         show_progress : bool, optional
-            If `True`, show a progress bar during the iteration over the
-            scans. Default: False
+            If ``True``, show a progress bar during the iteration over the
+            scans. Default: ``False``
 
         Notes
         -----
-        See `~pyvisgen.simulation.visibility.Polarisation` and
-        `~pyvisgen.simulation.visibility.Polarisation.rand_polarisation_field`
-        for more information on the keyword arguments in `pol_kwargs`
-        and `field_kwargs`, respectively.
+        See :class:`~pyvisgen.simulation.Polarisation` and
+        :class:`~pyvisgen.simulation.Polarisation.rand_polarisation_field`
+        for more information on the keyword arguments in ``pol_kwargs``
+        and ``field_kwargs``, respectively.
         """
         self.ra = torch.tensor(src_ra).double()
         self.dec = torch.tensor(src_dec).double()
@@ -396,6 +502,10 @@ def calc_time_steps(self):
         return time, time.mjd * (60 * 60 * 24)
 
     def calc_dense_baselines(self):
+        """Calculates the baselines of a densely-built
+        antenna array, which would provide full coverage of the
+        uv space.
+        """
         N = self.img_size
         fov = self.fov * pi / (3600 * 180)
         delta = fov ** (-1) * c.value / self.ref_frequency
@@ -437,9 +547,11 @@ def calc_dense_baselines(self):
         )
 
     def calc_baselines(self):
-        """Initializes Baselines dataclass object and
-        calls self.get_baselines to compute the contents of
-        the Baselines dataclass.
+        """Initializes :class:`~pyvisgen.simulation.Baselines`
+        dataclass object and calls
+        :py:func:`~pyvisgen.simulation.Observation.get_baselines`
+        to compute the contents of the :class:`~pyvisgen.simulation.Baselines`
+        dataclass.
         """
         self.baselines = Baselines(
             torch.tensor([]),  # st1
@@ -534,7 +646,22 @@ def get_baselines(self, times):
 
         return baselines
 
-    def calc_ref_elev(self, time=None):
+    def calc_ref_elev(self, time=None) -> tuple:
+        """Calculates the station elevations for given
+        time steps.
+
+        Parameters
+        ----------
+        time : array_like or None, optional
+            Array containing observation time steps.
+            Default: ``None``
+
+        Returns
+        -------
+        tuple
+            Tuple containing tensors of the Greenwich hour angle,
+            antenna-local hour angles, and the elevations.
+        """
         if time is None:
             time = self.times
         if time.shape == ():
@@ -592,7 +719,8 @@ def calc_feed_rotation(self, ha: Angle) -> Angle:
 
         .. math::
 
-            q = \atan\left(\frac{\sin h}{\cos\delta \tan\varphi - \sin\delta \cos h\right),
+            q = \atan\left(\frac{\sin h}{\cos\delta
+            \tan\varphi - \sin\delta \cos h\right),
 
         where $h$ is the local hour angle, $\varphi$ the geographical latitude
         of the observer, and $\delta$ the declination of the source.
@@ -612,30 +740,6 @@ def calc_feed_rotation(self, ha: Angle) -> Angle:
 
         return q
 
-    def calc_direction_cosines(self, ha, el_st, delta_x, delta_y, delta_z):
-        src_dec = torch.deg2rad(self.dec)
-        ha = torch.deg2rad(ha)
-
-        u = (torch.sin(ha) * delta_x + torch.cos(ha) * delta_y).reshape(-1)
-        v = (
-            -torch.sin(src_dec) * torch.cos(ha) * delta_x
-            + torch.sin(src_dec) * torch.sin(ha) * delta_y
-            + torch.cos(src_dec) * delta_z
-        ).reshape(-1)
-        w = (
-            torch.cos(src_dec) * torch.cos(ha) * delta_x
-            - torch.cos(src_dec) * torch.sin(ha) * delta_y
-            + torch.sin(src_dec) * delta_z
-        ).reshape(-1)
-
-        if not (u.shape == v.shape == w.shape):
-            raise ValueError(
-                "Expected u, v, and w to have the same shapes: "
-                f"{u.shape}, {v.shape}, {w.shape}"
-            )
-
-        return u, v, w
-
     def create_rd_grid(self):
         """Calculates RA and Dec values for a given fov around a source position
 

pyvisgen/simulation/visibility.py

@@ -10,6 +10,24 @@
 
 @dataclass
 class Visibilities:
+    """Visibilities dataclass.
+
+    Attributes
+    ----------
+    V_11 : :func:`~torch.tensor`
+    V_22 : :func:`~torch.tensor`
+    V_12 : :func:`~torch.tensor`
+    V_21 : :func:`~torch.tensor`
+    num : :func:`~torch.tensor`
+    base_num : :func:`~torch.tensor`
+    u : :func:`~torch.tensor`
+    v : :func:`~torch.tensor`
+    w  : :func:`~torch.tensor`
+    date : :func:`~torch.tensor`
+    linear_dop : :func:`~torch.tensor`
+    circular_dop : :func:`~torch.tensor`
+    """
+
     V_11: torch.tensor
     V_22: torch.tensor
     V_12: torch.tensor