Fixed restore run forced measurements associations

CHANGELOG.md

@@ -39,6 +39,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 #### Fixed
 
+- Fixed forced measurements being removed from associations during the restore run process [#600](https://github.com/askap-vast/vast-pipeline/pull/600).
 - Fixed measurement FITS cutout bug [#588](https://github.com/askap-vast/vast-pipeline/pull/588).
 - Fixed removal of image and sky region objects when a run is deleted [#585](https://github.com/askap-vast/vast-pipeline/pull/585).
 - Fixed testing pandas equal deprecation warning [#580](https://github.com/askap-vast/vast-pipeline/pull/580).
@@ -56,6 +57,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 #### List of PRs
 
+- [#600](https://github.com/askap-vast/vast-pipeline/pull/600): fix: Fixed restore run forced measurements associations.
 - [#590](https://github.com/askap-vast/vast-pipeline/pull/590): fix: Remove MeasurementPair model.
 - [#589](https://github.com/askap-vast/vast-pipeline/pull/589): fix: expose django-q timeout and retry to env vars.
 - [#588](https://github.com/askap-vast/vast-pipeline/pull/588): fix: change cutout endpoint to use measurement ID.

vast_pipeline/management/commands/restorepiperun.py

@@ -125,8 +125,6 @@ def restore_pipe(p_run: Run, bak_files: Dict[str, str], prev_config: PipelineCon
                     ' Cannot restore pipeline run.'
                 )
 
-            del meas
-
     logger.info("Restoring '%s' from backup parquet files.", p_run.name)
 
     # Delete any new sources
@@ -220,15 +218,20 @@ def restore_pipe(p_run: Run, bak_files: Dict[str, str], prev_config: PipelineCon
         with transaction.atomic():
             p_run.image_set.remove(*images_to_remove)
 
-    # load image meas
-    meas = pd.concat(
-        [pd.read_parquet(
-            i, columns=['id']
-        ) for i in prev_images['measurements_path']]
+    # load old associations to remove all new assoc
+    bak_assoc = pd.read_parquet(
+        bak_files['associations'],
+        columns=['source_id', 'meas_id']
     )
 
-    association_criteria_1 = Q(source_id__in=bak_sources['id'].to_numpy())
-    association_criteria_2 = ~Q(meas_id__in=meas['id'].to_numpy())
+    # get unique source and meas id values in the previous run
+    bak_source_ids = bak_assoc['source_id'].unique()
+    bak_meas_ids = bak_assoc['meas_id'].unique()
+
+    # create query to only obtain associations that are not part of the
+    # previous run
+    association_criteria_1 = Q(source_id__in=bak_source_ids)
+    association_criteria_2 = ~Q(meas_id__in=bak_meas_ids)
     associations_to_delete = Association.objects.filter(
         association_criteria_1 and association_criteria_2
     )
@@ -246,9 +249,11 @@ def restore_pipe(p_run: Run, bak_files: Dict[str, str], prev_config: PipelineCon
     logger.info('Restoring run metrics.')
     p_run.n_images = prev_images.shape[0]
     p_run.n_sources = bak_sources.shape[0]
-    p_run.n_selavy_measurements = meas.shape[0]
     if monitor:
+        p_run.n_selavy_measurements = meas.shape[0]
         p_run.n_forced_measurements = forced_meas.shape[0]
+    else:
+        p_run.n_selavy_measurements = bak_meas_ids.shape[0]
 
     with transaction.atomic():
         p_run.save()

vast_pipeline/tests/test_regression/compare_runs.py

@@ -4,7 +4,12 @@
 from django.test import TestCase
 
 
-def test_inc_assoc(testcase: TestCase, ass_add: pd.DataFrame, ass_backup: pd.DataFrame):
+def test_inc_assoc(
+    testcase: TestCase,
+    ass_add: pd.DataFrame,
+    ass_backup: pd.DataFrame,
+    must_be_equal: bool = False
+):
     '''
     Test that the number of associations increased or equal with added
     images.
@@ -17,9 +22,13 @@ def test_inc_assoc(testcase: TestCase, ass_add: pd.DataFrame, ass_backup: pd.Dat
         Associations after images were added.
     ass_backup : pd.DataFrame
         Associations before images were added.
+    must_be_equal: bool
+        The associations being compared must be equal in length to assert True.
     '''
-
-    testcase.assertTrue(len(ass_add) >= len(ass_backup))
+    if not must_be_equal:
+        testcase.assertTrue(len(ass_add) >= len(ass_backup))
+    else:
+        testcase.assertEqual(len(ass_add), len(ass_backup))
 
 def test_update_source(
     testcase: TestCase, sources_backup: pd.DataFrame, sources_backup_db: pd.DataFrame,

vast_pipeline/tests/test_regression/test_restore.py

@@ -35,8 +35,8 @@ def setUpTestData(self):
         '''
         Set up directories to test data, run the pipeline, and read the files.
         '''
-        base_path = 'normal-basic'
-        compare_path = 'restore-basic'
+        base_path = 'normal-basic-forced'
+        compare_path = 'restore-basic-forced'
         self.base_run = os.path.join(
             s.PIPELINE_WORKING_DIR, base_path
         )
@@ -109,7 +109,9 @@ def test_inc_assoc(self):
         '''
         See documentation for test_inc_assoc in compare_runs.
         '''
-        compare_runs.test_inc_assoc(self, self.ass_compare, self.ass_base)
+        compare_runs.test_inc_assoc(
+            self, self.ass_compare, self.ass_base, must_be_equal=True
+        )
 
     def test_update_source(self):
         '''