-
Notifications
You must be signed in to change notification settings - Fork 11
/
precomposed.py
109 lines (85 loc) · 3.83 KB
/
precomposed.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import typing
import torch
from mpol import fourier, images
from mpol.coordinates import GridCoords
class GriddedNet(torch.nn.Module):
r"""
.. note::
This module is provided as a starting point. However, we recommend
that you *don't get too comfortable using it* and instead write your own
(custom) modules following PyTorch idioms, potentially
using the source of this routine as a reference point. Using
the torch module system directly is *much more powerful* and expressive.
A basic but functional network for RML imaging. Designed to optimize a model image
using the entirety of the dataset in a :class:`mpol.datasets.GriddedDataset`
(i.e., gradient descent). For stochastic gradient descent (SGD), where the model
is only seeing a fraction of the dataset with each iteration, we recommend defining
your own module in your analysis code, following the 'Getting Started' guide.
.. mermaid:: ../_static/mmd/src/GriddedNet.mmd
Parameters
----------
coords : :class:`mpol.coordinates.GridCoords`
nchan : int
the number of channels in the base cube. Default = 1.
base_cube : :class:`mpol.images.BaseCube` or ``None``
a pre-packed base cube to initialize the model with. If
None, assumes ``torch.zeros``.
After the object is initialized, instance variables can be accessed, for example
:ivar bcube: the :class:`~mpol.images.BaseCube` instance
:ivar icube: the :class:`~mpol.images.ImageCube` instance
:ivar fcube: the :class:`~mpol.fourier.FourierCube` instance
For example, you'll likely want to access the ``self.icube.sky_model``
at some point.
"""
def __init__(
self,
coords: GridCoords,
nchan: int = 1,
base_cube: typing.Optional[torch.Tensor] = None,
) -> None:
super().__init__()
self.coords = coords
self.nchan = nchan
self.bcube = images.BaseCube(
coords=self.coords, nchan=self.nchan, base_cube=base_cube
)
self.conv_layer = images.HannConvCube(nchan=self.nchan)
self.icube = images.ImageCube(coords=self.coords, nchan=self.nchan)
self.fcube = fourier.FourierCube(coords=self.coords)
self.nufft = fourier.NuFFT(coords=self.coords, nchan=self.nchan)
def forward(self) -> torch.Tensor:
r"""
Feed forward to calculate the model visibilities. In this step, a
:class:`~mpol.images.BaseCube` is fed to a :class:`~mpol.images.HannConvCube`
is fed to a :class:`~mpol.images.ImageCube` is fed to a
:class:`~mpol.fourier.FourierCube` to produce the visibility cube.
Returns
-------
1D complex torch tensor of model visibilities.
"""
x = self.bcube()
x = self.conv_layer(x)
x = self.icube(x)
vis: torch.Tensor = self.fcube(x)
return vis
def predict_loose_visibilities(
self, uu: torch.Tensor, vv: torch.Tensor
) -> torch.Tensor:
r"""
Use the :class:`mpol.fourier.NuFFT` to calculate loose model visibilities from
the cube stored to ``self.icube.packed_cube``.
Parameters
----------
uu : :class:`torch.Tensor` of `class:`torch.double`
array of u spatial frequency coordinates,
not including Hermitian pairs. Units of [:math:`\mathrm{k}\lambda`]
vv : :class:`torch.Tensor` of `class:`torch.double`
array of v spatial frequency coordinates,
not including Hermitian pairs. Units of [:math:`\mathrm{k}\lambda`]
Returns
-------
:class:`torch.Tensor` of `class:`torch.complex128`
model visibilities corresponding to ``uu`` and ``vv`` locations.
"""
model_vis: torch.Tensor = self.nufft(self.icube.packed_cube, uu, vv)
return model_vis