gridder.py 4.22 KB
Newer Older
Philipp Arras's avatar
Philipp Arras committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2019 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.

Martin Reinecke's avatar
Martin Reinecke committed
18
19
import numpy as np

Philipp Arras's avatar
Philipp Arras committed
20
21
22
23
24
25
26
27
from ..domain_tuple import DomainTuple
from ..domains.rg_space import RGSpace
from ..domains.unstructured_domain import UnstructuredDomain
from ..operators.linear_operator import LinearOperator
from ..sugar import from_global_data, makeDomain


class GridderMaker(object):
28
    def __init__(self, dirty_domain, uv, eps=2e-13):
Martin Reinecke's avatar
Martin Reinecke committed
29
30
        import nifty_gridder
        dirty_domain = makeDomain(dirty_domain)
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
31
32
        if (len(dirty_domain) != 1 or not isinstance(dirty_domain[0], RGSpace)
                or not len(dirty_domain.shape) == 2):
Martin Reinecke's avatar
Martin Reinecke committed
33
            raise ValueError("need dirty_domain with exactly one 2D RGSpace")
34
35
36
37
        if uv.ndim != 2:
            raise ValueError("uv must be a 2D array")
        if uv.shape[1] != 2:
            raise ValueError("second dimension of uv must have length 2")
38
        dstx, dsty = dirty_domain[0].distances
39
        # wasteful hack to adjust to shape required by nifty_gridder
Martin Reinecke's avatar
Martin Reinecke committed
40
41
42
        uvw = np.empty((uv.shape[0], 3), dtype=np.float64)
        uvw[:, 0:2] = uv
        uvw[:, 2] = 0.
Martin Reinecke's avatar
typo    
Martin Reinecke committed
43
        # Scale uv such that 0<uv<=1 which is assumed by nifty_gridder
Martin Reinecke's avatar
Martin Reinecke committed
44
45
        uvw[:, 0] = uvw[:, 0]*dstx
        uvw[:, 1] = uvw[:, 1]*dsty
46
47
        speedOfLight = 299792458.
        bl = nifty_gridder.Baselines(uvw, np.array([speedOfLight]))
Martin Reinecke's avatar
Martin Reinecke committed
48
        nxdirty, nydirty = dirty_domain.shape
49
50
51
        gconf = nifty_gridder.GridderConfig(nxdirty, nydirty, eps, 1., 1.)
        nu, nv = gconf.Nu(), gconf.Nv()
        self._idx = nifty_gridder.getIndices(
Martin Reinecke's avatar
Martin Reinecke committed
52
            bl, gconf, np.zeros((uv.shape[0], 1), dtype=np.bool))
Philipp Arras's avatar
Philipp Arras committed
53
        self._bl = bl
Martin Reinecke's avatar
Martin Reinecke committed
54

Martin Reinecke's avatar
Martin Reinecke committed
55
56
        du, dv = 1./(nu*dstx), 1./(nv*dsty)
        grid_domain = RGSpace([nu, nv], distances=[du, dv], harmonic=True)
Martin Reinecke's avatar
Martin Reinecke committed
57
58

        self._rest = _RestOperator(dirty_domain, grid_domain, gconf)
Philipp Arras's avatar
Philipp Arras committed
59
        self._gridder = RadioGridder(grid_domain, bl, gconf, self._idx)
Martin Reinecke's avatar
Martin Reinecke committed
60
61
62

    def getGridder(self):
        return self._gridder
Philipp Arras's avatar
Philipp Arras committed
63
64
65
66

    def getRest(self):
        return self._rest

Martin Reinecke's avatar
Martin Reinecke committed
67
68
    def getFull(self):
        return self.getRest() @ self._gridder
Philipp Arras's avatar
Philipp Arras committed
69

Philipp Arras's avatar
Philipp Arras committed
70
71
72
    def ms2vis(self, x):
        return self._bl.ms2vis(x, self._idx)

Philipp Arras's avatar
Philipp Arras committed
73
74

class _RestOperator(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
75
76
77
78
    def __init__(self, dirty_domain, grid_domain, gconf):
        self._domain = makeDomain(grid_domain)
        self._target = makeDomain(dirty_domain)
        self._gconf = gconf
Philipp Arras's avatar
Philipp Arras committed
79
80
81
82
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
83
        res = x.to_global_data()
Philipp Arras's avatar
Philipp Arras committed
84
        if mode == self.TIMES:
85
            res = self._gconf.grid2dirty(res)
Philipp Arras's avatar
Philipp Arras committed
86
        else:
87
            res = self._gconf.dirty2grid(res)
Philipp Arras's avatar
Philipp Arras committed
88
89
90
91
        return from_global_data(self._tgt(mode), res)


class RadioGridder(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
92
    def __init__(self, grid_domain, bl, gconf, idx):
Martin Reinecke's avatar
Martin Reinecke committed
93
        self._domain = DomainTuple.make(
94
            UnstructuredDomain((bl.Nrows())))
Martin Reinecke's avatar
Martin Reinecke committed
95
96
97
98
        self._target = DomainTuple.make(grid_domain)
        self._bl = bl
        self._gconf = gconf
        self._idx = idx
Philipp Arras's avatar
Philipp Arras committed
99
100
101
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
Martin Reinecke's avatar
Martin Reinecke committed
102
        import nifty_gridder
Philipp Arras's avatar
Philipp Arras committed
103
104
        self._check_input(x, mode)
        if mode == self.TIMES:
105
            x = self._bl.ms2vis(x.to_global_data().reshape((-1, 1)), self._idx)
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
106
            res = nifty_gridder.vis2grid(self._bl, self._gconf, self._idx, x)
Philipp Arras's avatar
Philipp Arras committed
107
        else:
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
108
109
            res = nifty_gridder.grid2vis(self._bl, self._gconf, self._idx,
                                         x.to_global_data())
110
            res = self._bl.vis2ms(res, self._idx).reshape((-1,))
Philipp Arras's avatar
Philipp Arras committed
111
        return from_global_data(self._tgt(mode), res)