2222import  numpy  as  np 
2323import  pandas  as  pd 
2424import  unittest 
25- 
25+ import   lsst . afw . image   as   afwImage 
2626import  lsst .geom  as  geom 
2727import  lsst .utils .tests 
2828
@@ -46,20 +46,30 @@ def setUp(self):
4646        self .diaSources  =  pd .DataFrame (data = [
4747            {"ra" : 0.04 * idx  +  scatter * rng .uniform (- 1 , 1 ),
4848             "dec" : 0.04 * idx  +  scatter * rng .uniform (- 1 , 1 ),
49-              "diaSourceId" : idx  +  1  +  self .nObjects , "diaObjectId" : 0 }
49+              "diaSourceId" : idx  +  1  +  self .nObjects , "diaObjectId" : 0 ,  "trailLength" :  5.5 * idx }
5050            for  idx  in  range (self .nSources )])
5151        self .diaSourceZeroScatter  =  pd .DataFrame (data = [
5252            {"ra" : 0.04 * idx ,
5353             "dec" : 0.04 * idx ,
54-              "diaSourceId" : idx  +  1  +  self .nObjects , "diaObjectId" : 0 }
54+              "diaSourceId" : idx  +  1  +  self .nObjects , "diaObjectId" : 0 ,  "trailLength" :  5.5 * idx }
5555            for  idx  in  range (self .nSources )])
56+         visitInfo  =  afwImage .VisitInfo (exposureTime = 30.0 ,
57+                                        )
58+         exposureInfo  =  afwImage .ExposureInfo ()
59+         exposureInfo .setVisitInfo (visitInfo )
60+         maskedImage  =  afwImage .MaskedImageF (lsst .geom .Extent2I (64 , 64 ))
61+         self .exposure  =  afwImage .ExposureF (maskedImage , exposureInfo )
5662
5763    def  test_run (self ):
5864        """Test the full task by associating a set of diaSources to 
5965        existing diaObjects. 
6066        """ 
61-         assocTask  =  AssociationTask ()
62-         results  =  assocTask .run (self .diaSources , self .diaObjects )
67+ 
68+         config  =  AssociationTask .ConfigClass ()
69+         config .doTrailedSourceFilter  =  False 
70+         assocTask  =  AssociationTask (config = config )
71+ 
72+         results  =  assocTask .run (self .diaSources , self .diaObjects , self .exposure )
6373
6474        self .assertEqual (results .nUpdatedDiaObjects , len (self .diaObjects ) -  1 )
6575        self .assertEqual (results .nUnassociatedDiaObjects , 1 )
@@ -75,13 +85,36 @@ def test_run(self):
7585                [0 ]):
7686            self .assertEqual (test_obj_id , expected_obj_id )
7787
88+     def  test_run_trailed_sources (self ):
89+         """Test the full task by associating a set of diaSources to 
90+         existing diaObjects when trailed sources are filtered. 
91+         """ 
92+ 
93+         assocTask  =  AssociationTask ()
94+ 
95+         results  =  assocTask .run (self .diaSources , self .diaObjects , self .exposure )
96+ 
97+         self .assertEqual (results .nUpdatedDiaObjects , len (self .diaObjects ) -  3 )
98+         self .assertEqual (results .nUnassociatedDiaObjects , 3 )
99+         self .assertEqual (len (results .matchedDiaSources ),
100+                          len (self .diaObjects ) -  3 )
101+         self .assertEqual (len (results .unAssocDiaSources ), 1 )
102+         for  test_obj_id , expected_obj_id  in  zip (
103+                 results .matchedDiaSources ["diaObjectId" ].to_numpy (),
104+                 [1 , 2 , 3 , 4 ]):
105+             self .assertEqual (test_obj_id , expected_obj_id )
106+         for  test_obj_id , expected_obj_id  in  zip (
107+                 results .unAssocDiaSources ["diaObjectId" ].to_numpy (),
108+                 [0 ]):
109+             self .assertEqual (test_obj_id , expected_obj_id )
110+ 
78111    def  test_run_no_existing_objects (self ):
79112        """Test the run method with a completely empty database. 
80113        """ 
81114        assocTask  =  AssociationTask ()
82115        results  =  assocTask .run (
83116            self .diaSources ,
84-             pd .DataFrame (columns = ["ra" , "dec" , "diaObjectId" ]) )
117+             pd .DataFrame (columns = ["ra" , "dec" , "diaObjectId" ,  "trailLength" ]),  self . exposure )
85118        self .assertEqual (results .nUpdatedDiaObjects , 0 )
86119        self .assertEqual (results .nUnassociatedDiaObjects , 0 )
87120        self .assertEqual (len (results .matchedDiaSources ), 0 )
0 commit comments