Skip to content

Commit

Permalink
Add dataclass attributes docstring, add Polarisation class to __init__
Browse files Browse the repository at this point in the history
  • Loading branch information
@aknierim
aknierim committed Feb 1, 2025
1 parent 280c76a commit 69e6171
Show file tree
Hide file tree
Showing 3 changed files with 168 additions and 45 deletions.
3 changes: 2 additions & 1 deletion pyvisgen/simulation/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
192 changes: 148 additions & 44 deletions pyvisgen/simulation/observation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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: `{
Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 == ():
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down
18 changes: 18 additions & 0 deletions pyvisgen/simulation/visibility.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 69e6171

Please sign in to comment.