Polarisation

docs/changes/39.bugfix.rst

@@ -0,0 +1,4 @@
+- Fix gridding in `pyvisgen.simulation.observation.Observation` methods `create_rd_grid` and `create_lm_grid`
+  methods resulting in rotated images
+- Fix `pyvisgen.simulation.observation.ValidBaselineSubset` dataclass field order
+- Fix test failing because of api change

docs/changes/39.feature.rst

@@ -0,0 +1,4 @@
+- Add class `Polarisation` to `pyvisgen.simulation.visibility` that is called in `vis_loop`
+  - Add linear, circular, and no polarisation options
+- Update `pyvisgen.simulation.visibility.Visibilities` dataclass to also store polarisation degree tensors
+- Add keyword arguments for polarisation simulation to `pyvisgen.simulation.observation.Observation` class

docs/changes/39.maintenance.rst

@@ -0,0 +1,6 @@
+- Change pyvisgen.simulation.visibility.Visibilities dataclass component names from stokes components (I , Q, U, and V) 
+  to visibilities constructed from the stokes components (`V_11`, `V_22`, `V_12`, `V_21`)
+- Change indices for stokes components according to AIPS Memo 114
+  - Indices will be set automatically depending on simulated polarisation
+- Update comment strings in FITS files
+- Update docstrings accordingly in `pyvisgen.simulation.visibility.vis_loop` and `pyvisgen.simulation.observation.Observation`

pyvisgen/fits/writer.py

@@ -42,9 +42,17 @@
     freq_d = obs.bandwidths[0].cpu().numpy().item()
 
     ws = wcs.WCS(naxis=7)
+
+    crval_stokes = -1
+    stokes_comment = "-1=RR, -2=LL, -3=RL, -4=LR"
+    if obs.polarisation == "linear":
+        crval_stokes = -5
+        stokes_comment = "-5=XX, -6=YY, -7=XY, -8=YX"
+        stokes_comment += " or -5=VV, -6=HH, -7=VH, -8=HV"
+
     ws.wcs.crpix = [1, 1, 1, 1, 1, 1, 1]
     ws.wcs.cdelt = np.array([1, 1, -1, freq_d, 1, 1, 1])
-    ws.wcs.crval = [1, 1, -1, freq, 1, ra, dec]
+    ws.wcs.crval = [1, 1, crval_stokes, freq, 1, ra, dec]
     ws.wcs.ctype = ["", "COMPLEX", "STOKES", "FREQ", "IF", "RA", "DEC"]
     h = ws.to_header()
 
@@ -79,6 +87,8 @@
         hdu_vis.header["PZERO" + str(i + 1)] = parbzeros[i]
 
     # add comments
+    hdu_vis.header.comments["CTYPE2"] = "1=real, 2=imag, 3=weight"
+    hdu_vis.header.comments["CTYPE3"] = stokes_comment
     hdu_vis.header.comments["PTYPE1"] = "u baseline coordinate in light seconds"
     hdu_vis.header.comments["PTYPE2"] = "v baseline coordinate in light seconds"
     hdu_vis.header.comments["PTYPE3"] = "w baseline coordinate in light seconds"

pyvisgen/simulation/observation.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, fields
+from datetime import datetime
 from math import pi
 
 import astropy.constants as const
@@ -155,24 +156,104 @@ def _lexsort(self, a, dim=-1):
 class Observation:
     def __init__(
         self,
-        src_ra,
-        src_dec,
-        start_time,
-        scan_duration,
-        num_scans,
-        scan_separation,
-        integration_time,
-        ref_frequency,
-        frequency_offsets,
-        bandwidths,
-        fov,
-        image_size,
-        array_layout,
-        corrupted,
-        device,
-        dense=False,
-        sensitivity_cut=1e-6,
-    ):
+        src_ra: float,
+        src_dec: float,
+        start_time: datetime,
+        scan_duration: int,
+        num_scans: int,
+        scan_separation: int,
+        integration_time: int,
+        ref_frequency: float,
+        frequency_offsets: list,
+        bandwidths: list,
+        fov: float,
+        image_size: int,
+        array_layout: str,
+        corrupted: bool,
+        device: str,
+        dense: bool = False,
+        sensitivity_cut: float = 1e-6,
+        polarisation: str = None,
+        pol_kwargs: dict = {
+            "delta": 0,
+            "amp_ratio": 0.5,
+            "random_state": 42,
+        },
+        field_kwargs: dict = {
+            "order": [1, 1],
+            "scale": [0, 1],
+            "threshold": None,
+            "random_state": 42,
+        },
+    ) -> None:
+        """Sets up the observation class.
+
+        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 `~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
+            }`
+
+        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.
+        """
         self.ra = torch.tensor(src_ra).double()
         self.dec = torch.tensor(src_dec).double()
 
@@ -229,6 +310,11 @@ def __init__(
         self.rd = self.create_rd_grid()
         self.lm = self.create_lm_grid()
 
+        # polarisation
+        self.polarisation = polarisation
+        self.pol_kwargs = pol_kwargs
+        self.field_kwargs = field_kwargs
+
     def calc_dense_baselines(self):
         N = self.img_size
         fov = self.fov * pi / (3600 * 180)