Adding suggested ANC-related changes.

src/hivpy/hiv_testing.py

@@ -64,11 +64,11 @@ def update_hiv_testing(self, pop, time_step: timedelta):
         COVID disruption is factored in.
         """
         # mark people for testing
-        # hiv symptomatic > non hiv symptomatic > vmmc > anc > self testing > general > prep
+        # anc > hiv symptomatic > non hiv symptomatic > vmmc > self testing > general > prep
+        self.test_mark_anc(pop, time_step)
         self.test_mark_hiv_symptomatic(pop)
         self.test_mark_non_hiv_symptomatic(pop)
         self.test_mark_vmmc(pop, time_step)
-        self.test_mark_anc(pop, time_step)
         self.test_mark_general_pop(pop)
 
         # apply testing to marked population
@@ -152,7 +152,7 @@ def test_mark_non_hiv_symptomatic(self, pop):
            & (not (self.covid_disrup_affected | self.testing_disrup_covid))):
 
             # undiagnosed (last time step) population not scheduled for testing
-            not_diag_tested_pop = pop.get_sub_pop([(pop.get_correct_column(col.HIV_DIAGNOSED, dt=1), op.eq, False),
+            not_diag_tested_pop = pop.get_sub_pop([(col.HIV_DIAGNOSED, op.eq, False),
                                                    (col.TEST_MARK, op.eq, False)])
 
             if len(not_diag_tested_pop) > 0:
@@ -209,8 +209,6 @@ def test_mark_anc(self, pop, time_step):
                                                    (col.LAST_PREGNANCY_DATE, op.le, pop.date
                                                     - timedelta(days=270))])
             self.update_sub_pop_test_mark(pop, third_trimester_pop, self.prob_anc_test_trim3)
-            # remove from antenatal care
-            pop.set_present_variable(col.ANC, False, third_trimester_pop)
 
             # get post-delivery population tested during the previous time step
             post_delivery_pop = pop.get_sub_pop([(col.HIV_DIAGNOSED, op.eq, False),

src/hivpy/population.py

@@ -289,6 +289,9 @@ def evolve(self, time_step: timedelta):
             if (n_deaths and self.apply_death):
                 self.drop_from_population(HIV_deaths)
 
+        # Some population cleanup
+        self.pregnancy.reset_anc_at_birth(self)
+
         # Apply non-hiv deaths
         non_HIV_deaths = self.demographics.determine_deaths(self, time_step)
         n_deaths = n_deaths + sum(non_HIV_deaths)

src/hivpy/pregnancy.py

@@ -227,6 +227,19 @@ def update_want_no_children(self, pop: Population):
             # assign outcomes
             pop.set_present_variable(col.WANT_NO_CHILDREN, want_no_children, want_children_population)
 
+    def reset_anc_at_birth(self, pop: Population):
+        """
+        Reset the ANC status of everyone giving birth this time step
+        if they are currently in ANC.
+        """
+        # get population at the end of the third trimester
+        third_trimester_pop = pop.get_sub_pop([(col.ANC, op.eq, True),
+                                               (col.LAST_PREGNANCY_DATE, op.le, pop.date
+                                                - timedelta(days=270))])
+        if len(third_trimester_pop) > 0:
+            # remove from antenatal care
+            pop.set_present_variable(col.ANC, False, third_trimester_pop)
+
     def calc_prob_preg(self, age_group, ltp, stp, want_no_children):
         """
         Calculates the probability of getting pregnant for a group

src/tests/test_pregnancy.py

@@ -359,6 +359,7 @@ def update_anc_testing_outcomes(pop, time_step):
         pop.date += time_step
         pop.pregnancy.update_pregnancy(pop)
     update_anc_testing_outcomes(pop, time_step)
+    pop.pregnancy.reset_anc_at_birth(pop)
 
     # get stats
     no_tested = len(pop.get_sub_pop([(col.LAST_TEST_DATE, op.eq, pop.date)]))