Support parallactic angles computed by astropy as well as faster casa measures parallactic angle computation

montblanc/util/casa_parangles.cpp

@@ -0,0 +1,93 @@
+/*
+<%
+setup_pybind11(cfg)
+cfg['compiler_args'] = ['-std=c++11', '-fvisibility=hidden']
+cfg['libraries'] = ['casa_casa', 'casa_measures']
+%>
+*/
+
+#include <iostream>
+#include <vector>
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+
+#include <casacore/measures/Measures/MCDirection.h>
+#include <casacore/measures/Measures/MDirection.h>
+#include <casacore/measures/Measures/MPosition.h>
+#include <casacore/measures/Measures/MeasConvert.h>
+#include <casacore/measures/Measures/MeasTable.h>
+#include <casacore/measures/Measures/MEpoch.h>
+
+namespace py = pybind11;
+
+constexpr unsigned int flags = py::array::c_style | py::array::forcecast;
+
+template <typename FT>
+py::array_t<FT, flags> parallactic_angles(
+    py::array_t<FT, flags> times,
+    py::array_t<FT, flags> antenna_positions,
+    py::array_t<FT, flags> phase_centre)
+{
+    py::gil_scoped_release release;
+
+    int na = antenna_positions.shape(0);
+    int ntimes = times.shape(0);
+
+    // Result array
+    py::array_t<FT, flags> angles({ntimes, na});
+
+    std::vector<casa::MPosition> itrf_antenna_positions;
+    itrf_antenna_positions.reserve(na);
+
+    // Compute antenna positions in ITRF
+    for(int ant=0; ant<na; ++ant)
+    {
+        const FT * x = antenna_positions.data(ant, 0);
+        const FT * y = antenna_positions.data(ant, 1);
+        const FT * z = antenna_positions.data(ant, 2);
+
+        itrf_antenna_positions.push_back(casa::MPosition(
+                casa::MVPosition(*x, *y, *z),
+                casa::MPosition::ITRF));
+    }
+
+    // Direction towards zenith
+    casa::MVDirection base_zenith(0, M_PI/2);
+
+    // Direction towards phase centre
+    casa::MVDirection phase_dir(*phase_centre.data(0), *phase_centre.data(1));
+
+    // For each time
+    for(int time=0; time<ntimes; ++time)
+    {
+        // Create a frame for this timestemp
+        casa::MeasFrame frame(casa::MEpoch(
+            casa::Quantum<double>(*times.data(time), "s"),
+            casa::MEpoch::UTC));
+
+        // For each antenna
+        for(int ant=0; ant<na; ++ant)
+        {
+            // Set the frame's position to the antenna position
+            frame.set(itrf_antenna_positions[ant]);
+
+            // Direction to the zenith in this frame
+            casa::MDirection mzenith(base_zenith, casa::MDirection::Ref(
+                casa::MDirection::AZELGEO, frame));
+
+            // Compute parallactic angle of phase direction w.r.t zenith
+            casa::MeasConvert<casa::MDirection> convert(mzenith, casa::MDirection::J2000);
+            *angles.mutable_data(time, ant) = phase_dir.positionAngle(convert().getValue());
+        }
+    }
+
+    return angles;
+}
+
+PYBIND11_MODULE(casa_parangles, m) {
+    m.doc() = "auto-compiled c++ extension";
+
+    m.def("parallactic_angles", &parallactic_angles<float>, py::return_value_policy::move);
+    m.def("parallactic_angles", &parallactic_angles<double>, py::return_value_policy::move);
+}

montblanc/util/parallactic_angles.py

@@ -18,10 +18,16 @@
 # You should have received a copy of the GNU General Public License
 # along with this program; if not, see <http://www.gnu.org/licenses/>.
 
+import cppimport
 import numpy as np
+import pybind11
 
 import montblanc
 
+cppimport.set_quiet(False)
+mod = cppimport.imp("montblanc.util.casa_parangles")
+mod = mod.util.casa_parangles
+
 try:
     import pyrap.measures
 
@@ -54,7 +60,7 @@ def parallactic_angles(times, antenna_positions, field_centre):
 
     try:
         # Create direction measure for the zenith
-        zenith = pm.direction('AZEL','0deg','90deg')
+        zenith = pm.direction('AZELGEO','0deg','90deg')
     except AttributeError as e:
         if pm is None:
             raise ImportError("python-casacore import failed")
@@ -80,4 +86,112 @@ def parallactic_angles(times, antenna_positions, field_centre):
                 pm.posangle(fc_rad, zenith).get_value("rad")
                 for rp in reference_positions
             ]
-        for t in times])
+        for t in times])
+
+
+from astropy.coordinates import EarthLocation, SkyCoord, AltAz, CIRS, Angle
+from astropy.time import Time
+
+from astropy import units
+
+def _parallactic_angle_astropy(times, ap, fc):
+    """
+    Computes parallactic angles per timestep for the given
+    reference antenna position and field centre.
+
+    Arguments:
+        times: ndarray
+            Array of unique times with shape (ntime,),
+            obtained from TIME column of MS table
+        ap: ndarray of shape (na, 3)
+            Antenna positions, obtained from POSITION
+            column of MS ANTENNA sub-table
+        fc : ndarray of shape (2,)
+            Field centre, should be obtained from MS PHASE_DIR
+
+    Returns:
+        An array of parallactic angles per time-step
+
+    """
+    from astropy.coordinates import EarthLocation, SkyCoord, AltAz, CIRS
+    from astropy.time import Time
+    from astropy import units
+
+    # Convert from MJD second to MJD
+    times = Time(times / 86400.00, format='mjd', scale='utc')
+
+    ap = EarthLocation.from_geocentric(ap[:,0], ap[:,1], ap[:,2], unit='m')
+    fc = SkyCoord(ra=fc[0], dec=fc[1], unit=units.rad, frame='fk5')
+    pole = SkyCoord(ra=0, dec=90, unit=units.deg, frame='fk5')
+
+    cirs_frame = CIRS(obstime=times)
+    pole_cirs = pole.transform_to(cirs_frame)
+    fc_cirs = fc.transform_to(cirs_frame)
+
+    altaz_frame = AltAz(location=ap[None,:], obstime=times[:,None])
+    pole_altaz = pole_cirs[:,None].transform_to(altaz_frame)
+    fc_altaz = fc_cirs[:,None].transform_to(altaz_frame)
+    return fc_altaz.position_angle(pole_altaz)
+
+if __name__ == "__main__":
+    # 5s second types from 12h00 midday on 1st Feb
+    times = np.arange('2017-02-01T12:00', '2017-02-01T16:00', dtype='datetime64[5s]')
+    ftimes = times.astype('datetime64[s]').astype(np.float64)
+
+    # Westerbork antenna positions
+    antenna_positions = np.array([
+           [ 3828763.10544699,   442449.10566454,  5064923.00777   ],
+           [ 3828746.54957258,   442592.13950824,  5064923.00792   ],
+           [ 3828729.99081359,   442735.17696417,  5064923.00829   ],
+           [ 3828713.43109885,   442878.2118934 ,  5064923.00436   ],
+           [ 3828696.86994428,   443021.24917264,  5064923.00397   ],
+           [ 3828680.31391933,   443164.28596862,  5064923.00035   ],
+           [ 3828663.75159173,   443307.32138056,  5064923.00204   ],
+           [ 3828647.19342757,   443450.35604638,  5064923.0023    ],
+           [ 3828630.63486201,   443593.39226634,  5064922.99755   ],
+           [ 3828614.07606798,   443736.42941621,  5064923.        ],
+           [ 3828609.94224429,   443772.19450029,  5064922.99868   ],
+           [ 3828601.66208572,   443843.71178407,  5064922.99963   ],
+           [ 3828460.92418735,   445059.52053929,  5064922.99071   ],
+           [ 3828452.64716351,   445131.03744105,  5064922.98793   ]],
+                dtype=np.float64)
+
+    phase_centre = np.array([ 0.        ,  1.04719755], dtype=np.float64)
+
+    import time
+
+    t = time.clock()
+    cangles = mod.parallactic_angles(ftimes, antenna_positions, phase_centre)
+    print 'cangles dones in %f' % (time.clock() - t)
+
+    t = time.time()
+    pa_astro = _parallactic_angle_astropy(ftimes, antenna_positions, phase_centre)
+    print "pa_astropy done in ", time.time() - t
+
+    t = time.clock()
+    pangles = parallactic_angles(ftimes, antenna_positions, phase_centre)
+    print 'pangles dones in %f' % (time.clock() - t)
+
+    assert np.allclose(cangles, pangles)
+
+    import sys
+    sys.exit(0)
+
+    import time
+
+    t = time.time()
+    pa_astro = _parallactic_angle_astropy(ftimes, antenna_positions, phase_centre)
+    print "pa_astropy done in ", time.time() - t
+
+    t = time.time()
+    pa_casa = parallactic_angles(ftimes, antenna_positions, phase_centre)
+    print "pa_casa done in ", time.time() - t
+
+    pa_astro = Angle(pa_astro, unit=units.deg).wrap_at(180*units.deg)
+    pa_casa = Angle(pa_casa*units.rad, unit=units.deg).wrap_at(180*units.deg)
+
+    for a, c in zip(pa_astro.flat, pa_casa.flat):
+        print a, c, (a-c).wrap_at(180*units.deg)
+
+
+

setup.py

@@ -21,6 +21,7 @@
 import json
 import logging
 import os
+from os.path import join as pjoin
 import sys
 
 #==============
@@ -30,7 +31,7 @@
 from install.install_log import log
 
 mb_path = 'montblanc'
-mb_inc_path = os.path.join(mb_path, 'include')
+mb_inc_path = pjoin(mb_path, 'include')
 
 #===================
 # Detect readthedocs
@@ -142,10 +143,23 @@ def include_pkg_dirs():
             # OK, so everything starts with 'montblanc/'
             # Take everything after that ('include...') and
             # append a '/*.*' to it
-            pkg_dirs.append(os.path.join(root[l:], d, '*.*'))
+            pkg_dirs.append(pjoin(root[l:], d, '*.*'))
 
     return pkg_dirs
 
+
+def include_pkg_dirs():
+    mb_inc_path = pjoin("montblanc", "include")
+    l = len("montblanc") + len(os.sep)
+
+    exclude = set(['docs', '.git', '.svn'])
+
+    return [pjoin(path, d, '*.*')[l:] for path, dirs, files
+                    in os.walk(mb_inc_path, topdown=True)
+                    for d in dirs if d not in exclude]
+
+
+
 install_requires = [
     'attrdict >= 2.0.0',
     'attrs >= 16.3.0',
@@ -167,10 +181,11 @@ def include_pkg_dirs():
 else:
     # Add binary/C extension type packages
     install_requires += [
-        'astropy >= 1.3.0',
+        'astropy >= 2.0.2',
+        'cppimport >= 17.7.24',
         'cerberus >= 1.1',
         'numpy >= 1.11.3',
-        'numexpr >= 2.6.1',
+        'pybind11 >= 2.2.0',
         'python-casacore >= 2.1.2',
         'ruamel.yaml >= 0.15.22',
         "{} >= 1.3.0".format(tensorflow_package),
@@ -196,7 +211,7 @@ def include_pkg_dirs():
 from install.versioning import maintain_version
 
 setup(name='montblanc',
-    version=maintain_version(os.path.join('montblanc', 'version.py')),
+    version=maintain_version(pjoin('montblanc', 'version.py')),
     description='GPU-accelerated RIME implementations.',
     long_description=readme(),
     url='http://github.com/ska-sa/montblanc',
@@ -217,6 +232,5 @@ def include_pkg_dirs():
     license='GPL2',
     install_requires=install_requires,
     packages=find_packages(),
-    package_data={'montblanc': include_pkg_dirs()},
-    include_package_data=True,
+    package_data={'montblanc': include_pkg_dirs() + [pjoin("util", "*.cpp")] },
     zip_safe=False)