Fix aij flux calculation

stellarphot/differential_photometry/aij_rel_fluxes.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from astropy.coordinates import SkyCoord
+from astropy.table import Table
 import astropy.units as u
 
 __all__ = ['add_in_quadrature', 'calc_aij_relative_flux']
@@ -14,6 +15,7 @@ def add_in_quadrature(array):
 
 def calc_aij_relative_flux(star_data, comp_stars,
                            in_place=True, coord_column=None, 
+                           star_id_column='star_id',
                            flux_column_name='aperture_net_flux'):
     """
     Calculate AstroImageJ-style flux ratios.
@@ -30,14 +32,18 @@ def calc_aij_relative_flux(star_data, comp_stars,
     in_place : ``bool``, optional
         If ``True``, add new columns to input table. Otherwise, return
         new table with those columns added.
-        
+
     coord_column : ``str``, optional
         If provided, use this column to match comparison stars to coordinates.
         If not provided, the coordinates are generated with SkyCoord.
-        
+
     flux_column_name : ``str``, optional
-        If provided, use this column to find flux.  
-        If not provided, the column 'aperture_net_flux' is used. 
+        If provided, use this column to find flux.
+        If not provided, the column 'aperture_net_flux' is used.
+
+    star_id_column : ``str``, optional
+        Name of the column that provides a unique identifier for each
+        comparison star.
 
     Returns
     -------
@@ -61,6 +67,17 @@ def calc_aij_relative_flux(star_data, comp_stars,
     # Not sure this is really close enough for a good match...
     good = d2d < 1.2 * u.arcsec
 
+    check_for_bad = Table(data=[star_data[star_id_column].data, good],
+                          names=['star_id', 'good'])
+    check_for_bad = check_for_bad.group_by('star_id')
+    is_all_good = check_for_bad.groups.aggregate(np.all)
+
+    bad_comps = is_all_good['star_id'][~is_all_good['good']]
+
+    for comp in bad_comps:
+        this_comp = star_data[star_id_column] == comp
+        good[this_comp] = False
+
     error_column_name = 'noise-aij'
     # Calculate comp star counts for each time
 

stellarphot/differential_photometry/tests/test_aij_rel_fluxes.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from astropy.coordinates import SkyCoord
 from astropy.table import Table
 from astropy.time import Time
 import astropy.units as u
@@ -31,6 +32,7 @@ def _raw_photometry_table():
     star_dec = np.array([45.0] * n_stars) * u.degree
     fluxes = np.array([10000., 20000, 30000, 40000])
     errors = np.sqrt(fluxes) + 50
+    star_ids = np.arange(1, 5, dtype='int')
 
     # Stars 2, 3 and 4 will be the comparison stars
     comp_stars = np.array([0, 1, 1, 1])
@@ -47,22 +49,17 @@ def _raw_photometry_table():
 
     raw_table = Table(data=[np.sort(_repeat(times, n_stars)), _repeat(star_ra,  n_times),
                             _repeat(star_dec, n_times), _repeat(fluxes, n_times),
-                            _repeat(errors, n_times)],
-                      names=['date-obs', 'RA', 'Dec', 'aperture_net_flux', 'noise-aij'])
+                            _repeat(errors, n_times),
+                            _repeat(star_ids, n_times)],
+                      names=['date-obs', 'RA', 'Dec', 'aperture_net_flux',
+                             'noise-aij', 'star_id'])
 
     return expected_flux_ratios, expected_flux_error, raw_table, raw_table[1:4]
 
 
 @pytest.mark.parametrize('in_place', [True, False])
 def test_relative_flux_calculation(in_place):
     expected_flux, expected_error, input_table, comp_star = _raw_photometry_table()
-    # print(input_table)
-
-    # Try running the fluxes one exposure at a time
-    # for one_time in grouped_input.groups:
-    #     output = calc_aij_mags(one_time, comp_star)
-    #     #print(one_time, comp_star, output)
-    #     np.testing.assert_allclose(output, expected_flux)
 
     # Try doing it all at once
     n_times = len(np.unique(input_table['date-obs']))
@@ -80,3 +77,35 @@ def test_relative_flux_calculation(in_place):
         assert 'relative_flux' in input_table.colnames
     else:
         assert 'relative_flux' not in input_table.colnames
+
+
+def test_mislocated_comp_star():
+    expected_flux, expected_error, input_table, comp_star = \
+        _raw_photometry_table()
+    # "Jiggle" one of the stars by moving it by a few arcsec in one image.
+    # We'll do it in the last image.
+
+    # First, let's sort so the row we want to modify is the last one
+    input_table.sort(['date-obs', 'star_id'])
+    last_one = input_table[-1]
+    coord_inp = SkyCoord(ra=last_one['RA'], dec=last_one['Dec'], unit=u.degree)
+    coord_bad_ra = coord_inp.ra + 3 * u.arcsecond
+    input_table['RA'][-1] = coord_bad_ra.degree
+
+    output_table = calc_aij_relative_flux(input_table, comp_star,
+                                          in_place=False)
+
+    old_total_flux = comp_star['aperture_net_flux'].sum()
+    new_flux = old_total_flux - last_one['aperture_net_flux']
+    # This works for target stars, i.e. those never in comparison set
+    new_expected_flux = old_total_flux / new_flux * expected_flux
+
+    # Oh wow, this is terrible....
+    # Need to manually calculate for the only two that are still in comparison
+    new_expected_flux[1] = (comp_star['aperture_net_flux'][0] /
+                            comp_star['aperture_net_flux'][1])
+    new_expected_flux[2] = (comp_star['aperture_net_flux'][1] /
+                            comp_star['aperture_net_flux'][0])
+    new_expected_flux[3] = expected_flux[3]
+
+    np.testing.assert_allclose(new_expected_flux, output_table['relative_flux'][-4:])