get_trajectory for 1L2S models: proposed bug fix + unit tests

.gitignore

@@ -67,3 +67,4 @@ source/MulensModel/*.so
 
 # Coverage report:
 coverage.xml
+/source/MulensModel/htmlcov/

examples/example_17_1L2S_plotting.py

@@ -94,4 +94,9 @@ def generate_dataset(time, flux_1, flux_2, blend_flux, flux_err,
 my_event.plot_model(zorder=10, color='black')
 my_event.plot_data()
 
+# Plot the source trajectories
+plt.figure()
+plt.title('Model Source Trajectories')
+my_model.plot_trajectory()
+
 plt.show()

source/MulensModel/fitdata.py

@@ -647,13 +647,9 @@ def get_dataset_trajectory(self):
         if self.dataset.ephemerides_file is None:
             return self.model.get_trajectory(self.dataset.time)
         else:
-            kwargs_ = {
-                'times': self.dataset.time,
-                'parallax': self.model.get_parallax(),
-                'coords': self.model.coords,
-                'satellite_skycoord': self.dataset.satellite_skycoord}
+            return self.model.get_trajectory(
+                self.dataset.time, satellite_skycoord=self.dataset.satellite_skycoord)
 
-            return Trajectory(parameters=self.model.parameters, **kwargs_)
 
     def get_d_A_d_u_for_point_lens_model(self):
         """

source/MulensModel/model.py

@@ -747,7 +747,7 @@ def _plot_source_for_trajectory(self, trajectory, **kwargs):
             for (x, y) in zip(trajectory.x, trajectory.y):
                 axis.add_artist(plt.Circle((x, y), **kwargs))
 
-    def get_trajectory(self, times):
+    def get_trajectory(self, times, satellite_skycoord=None):
         """
         Get the source trajectory for the given set of times.
 
@@ -758,10 +758,20 @@ def get_trajectory(self, times):
         Returns : A `:py:class:`~MulensModel.trajectory.Trajectory` object.
 
         """
+        if satellite_skycoord is None:
+            satellite_skycoord = self.get_satellite_coords(times)
+
         kwargs_ = {
             'times': times, 'parallax': self._parallax, 'coords': self._coords,
-            'satellite_skycoord': self.get_satellite_coords(times)}
-        return Trajectory(parameters=self.parameters, **kwargs_)
+            'satellite_skycoord': satellite_skycoord}
+        if self.n_sources == 1:
+            return Trajectory(parameters=self.parameters, **kwargs_)
+        elif self.n_sources == 1:
+            trajectory_1 = Trajectory(
+                parameters=self.parameters.source_1_parameters, **kwargs_)
+            trajectory_2 = Trajectory(
+                parameters=self.parameters.source_2_parameters, **kwargs_)
+            return (trajectory_1, trajectory_2)
 
     def set_times(
             self, t_range=None, t_start=None, t_stop=None, dt=None,

source/MulensModel/tests/test_FitData.py

@@ -621,6 +621,31 @@ def test_photfmt_scaled_2(self):
         almost(res_errors, self.dataset.err_mag)
 
 
+def test_get_trajectory_1L2S_satellite_parallax():
+    """test parallax calculation with Spitzer data"""
+    model_parameters = {'t_0': 2456836.22, 'u_0': 0.922, 't_E': 22.87,
+                        'pi_E_N': -0.248, 'pi_E_E': 0.234,
+                        't_0_par': 2456836.2}
+    coords = "17:47:12.25 -21:22:58.2"
+
+    model_with_par = mm.Model(model_parameters, coords=coords)
+    model_with_par.parallax(satellite=True, earth_orbital=True,
+                            topocentric=False)
+
+    data_Spitzer = mm.MulensData(
+        file_name=SAMPLE_FILE_03, ephemerides_file=SAMPLE_FILE_03_EPH)
+
+    fit = mm.FitData(model=model_with_par, dataset=data_Spitzer)
+
+    ref_Spitzer = np.loadtxt(SAMPLE_FILE_03_REF, unpack=True)
+
+    trajectory = fit.get_dataset_trajectory()
+
+    ratio_x = trajectory.x / ref_Spitzer[6]
+    ratio_y = trajectory.x / ref_Spitzer[7]
+    np.testing.assert_almost_equal(ratio_x, [1.]*len(ratio_x), decimal=4)
+    np.testing.assert_almost_equal(ratio_y, [1.] * len(ratio_y), decimal=4)
+
 # Tests to add:
 #
 # test get_chi2_gradient(), chi2_gradient:

source/MulensModel/tests/test_Model.py

@@ -3,9 +3,16 @@
 from math import isclose
 import unittest
 from astropy import units as u
+import os.path
 
 import MulensModel as mm
 
+dir_2 = os.path.join(mm.DATA_PATH, 'unit_test_files')
+dir_3 = os.path.join(mm.DATA_PATH, 'ephemeris_files')
+SAMPLE_FILE_02_REF = os.path.join(dir_2, 'ob140939_OGLE_ref_v1.dat')  # HJD'
+SAMPLE_FILE_03_EPH = os.path.join(dir_3, 'Spitzer_ephemeris_01.dat')  # UTC
+SAMPLE_FILE_03_REF = os.path.join(dir_2, 'ob140939_Spitzer_ref_v1.dat')  # HJD'
+
 
 def test_n_lenses():
     """check n_lenses property"""
@@ -486,6 +493,97 @@ def test_repr():
     assert str(model) == expected
 
 
+class TestGetTrajectory(unittest.TestCase):
+
+    def setUp(self):
+        # Parallax model parameters
+        self.model_parameters_par = {'t_0': 2456836.22, 'u_0': 0.922, 't_E': 22.87,
+                            'pi_E_N': -0.248, 'pi_E_E': 0.234,
+                            't_0_par': 2456836.2}
+        self.coords = "17:47:12.25 -21:22:58.2"
+
+        self.model_with_par = mm.Model(
+            self.model_parameters_par, coords=self.coords)
+        self.model_with_par.parallax(satellite=True, earth_orbital=True,
+                                topocentric=False)
+
+        self.ref_OGLE = np.loadtxt(SAMPLE_FILE_02_REF, unpack=True)
+        self.times_OGLE = self.ref_OGLE[0]
+        self.ref_Spitzer = np.loadtxt(SAMPLE_FILE_03_REF, unpack=True)
+        self.ephemerides_file = SAMPLE_FILE_03_EPH
+        self.times_spz = self.ref_Spitzer[0]
+
+    def test_1L1S(self):
+        # straight-up trajectory for static point-lens model
+        t_0 = self.model_parameters_par['t_0']
+        u_0 = self.model_parameters_par['u_0']
+        t_E = self.model_parameters_par['t_E']
+        times = np.arange(t_0 - t_E, t_0 + t_E, 1.)
+
+        model = mm.Model({'t_0': t_0, 'u_0': u_0, 't_E': t_E})
+        trajectory = model.get_trajectory(times)
+        y = np.ones(len(times)) * u_0
+        x = (times - t_0) / t_E
+        ratio_x = trajectory.x / x
+        ratio_y = trajectory.y / y
+        np.testing.assert_almost_equal(ratio_x, [1.] * len(ratio_x), decimal=4)
+        np.testing.assert_almost_equal(ratio_y, [1.] * len(ratio_y), decimal=4)
+
+    def test_1L1S_annual_parallax(self):
+        # case with annual parallax
+        trajectory = self.model_with_par.get_trajectory(
+            self.times_OGLE)
+
+        ratio_x = trajectory.x / self.ref_OGLE[6]
+        ratio_y = trajectory.y / self.ref_OGLE[7]
+        np.testing.assert_almost_equal(ratio_x, [1.] * len(ratio_x), decimal=4)
+        np.testing.assert_almost_equal(ratio_y, [1.] * len(ratio_y), decimal=4)
+
+    def test_1L1S_satellite_parallax_1(self):
+        # case with satellite parallax (check test_Model_Parallax.py)
+        satellite_skycoord_obj = mm.SatelliteSkyCoord(
+            ephemerides_file=self.ephemerides_file)
+        satellite_skycoord = satellite_skycoord_obj.get_satellite_coords(
+            self.times_spz)
+        trajectory = self.model_with_par.get_trajectory(
+            self.times_spz, satellite_skycoord=satellite_skycoord)
+
+        ratio_x = trajectory.x / self.ref_Spitzer[6]
+        ratio_y = trajectory.y / self.ref_Spitzer[7]
+        np.testing.assert_almost_equal(ratio_x, [1.] * len(ratio_x), decimal=4)
+        np.testing.assert_almost_equal(ratio_y, [1.] * len(ratio_y), decimal=4)
+
+    def test_1L1S_satellite_parallax_2(self):
+        # case with satellite parallax (check test_Model_Parallax.py)
+        model_with_sat_par = mm.Model(
+            self.model_parameters_par, ephemerides_file=self.ephemerides_file,
+            coords=self.coords)
+        trajectory = model_with_sat_par.get_trajectory(self.times_spz)
+
+        ratio_x = trajectory.x / self.ref_Spitzer[6]
+        ratio_y = trajectory.y / self.ref_Spitzer[7]
+        np.testing.assert_almost_equal(ratio_x, [1.] * len(ratio_x), decimal=4)
+        np.testing.assert_almost_equal(ratio_y, [1.] * len(ratio_y), decimal=4)
+
+    def test_1L2S(self):
+        # Binary source trajectories
+        model = mm.Model({
+            't_0_1': 5000., 'u_0_1': 0.005, 'rho_1': 0.001,
+            't_0_2': 5100., 'u_0_2': 0.0003, 't_star_2': 0.03, 't_E': 25.})
+        model_1 = mm.Model(model.parameters.source_1_parameters)
+        model_2 = mm.Model(model.parameters.source_2_parameters)
+
+        time = np.linspace(4900., 5200., 4200)
+
+        (traj_1_1L2S, traj_2_1L2S) = model.get_trajectory(time)
+        traj_1_1L1S = model_1.get_trajectory(time)
+        np.testing.assert_equal(traj_1_1L2S.x == traj_1_1L1S.x)
+        np.testing.assert_equal(traj_1_1L2S.y == traj_1_1L1S.y)
+
+        traj_2_1L1S = model_2.get_trajectory(time)
+        np.testing.assert_equal(traj_2_1L2S.x == traj_2_1L1S.x)
+        np.testing.assert_equal(traj_2_1L2S.y == traj_2_1L1S.y)
+
 # Tests to Add:
 #
 # test get_trajectory:

source/MulensModel/tests/test_Trajectory.py

@@ -19,10 +19,10 @@ def test_coords_format():
     coords = [None, coords_txt, mm.Coordinates(coords_txt),
               SkyCoord(coords_txt, unit=(u.hourangle, u.deg))]
 
-    trajecotories = [mm.Trajectory(times, params, coords=c) for c in coords]
-    for trajectory in trajecotories:
-        assert np.all(trajectory.x == trajecotories[0].x)
-        assert np.all(trajectory.y == trajecotories[0].y)
+    trajectories = [mm.Trajectory(times, params, coords=c) for c in coords]
+    for trajectory in trajectories:
+        assert np.all(trajectory.x == trajectories[0].x)
+        assert np.all(trajectory.y == trajectories[0].y)
 
     parameters['pi_E_E'] = 0.1
     parameters['pi_E_N'] = -0.15
@@ -31,7 +31,7 @@ def test_coords_format():
     coords = coords[1:]
     p = {'earth_orbital': True, 'satellite': False, 'topocentric': False}
     kwargs = {'times': times, 'parameters': params, 'parallax': p}
-    trajecotories = [mm.Trajectory(coords=c, **kwargs) for c in coords]
-    for trajectory in trajecotories:
-        assert np.all(trajectory.x == trajecotories[0].x)
-        assert np.all(trajectory.y == trajecotories[0].y)
+    trajectories = [mm.Trajectory(coords=c, **kwargs) for c in coords]
+    for trajectory in trajectories:
+        assert np.all(trajectory.x == trajectories[0].x)
+        assert np.all(trajectory.y == trajectories[0].y)