Merge pull request #190 from arup-group/fs-optional-override

retain facility locations

pam/core.py

@@ -451,7 +451,8 @@ def combine(self, other, prefix=""):
         raise TypeError(
             f"Object for addition must be a Population Household or Person object, not {type(other)}")
 
-    def sample_locs(self, sampler, long_term_activities=None, joint_trips_prefix='escort_'):
+    def sample_locs(self, sampler, long_term_activities=None, joint_trips_prefix='escort_',
+                    location_override=True):
         """
         WIP Sample household plan locs using a sampler.
 
@@ -472,8 +473,10 @@ def sample_locs(self, sampler, long_term_activities=None, joint_trips_prefix='es
         persons and home activities are impacted.
 
         TODO - add method to all core classes
-        :params list long_term activities: a list of activities for which location is only assigned once (per zone)
-        :params str joint_trips_prefix: a purpose prefix used to identify escort/joint trips
+        :param list long_term activities: a list of activities for which location is only assigned once (per zone)
+        :param str joint_trips_prefix: a purpose prefix used to identify escort/joint trips
+        :param bool location_override: if False, the facility sampler will retain any 
+            already-existing locations in the population.
         """
         if long_term_activities is None:
             long_term_activities = variables.LONG_TERM_ACTIVITIES
@@ -496,13 +499,14 @@ def sample_locs(self, sampler, long_term_activities=None, joint_trips_prefix='es
                         target_act = act.act[(len(joint_trips_prefix)):]
                     else:
                         target_act = act.act
-
+
+                    # assign any unique locations
                     if (act.location.area, target_act) in unique_locations:
                         location = unique_locations[(
                             act.location.area, target_act)]
                         act.location = location
-
-                    else:
+                    # sample facility
+                    elif location_override or act.location.loc is None:
                         location = activity.Location(
                             area=act.location.area,
                             loc=sampler.sample(act.location.area, target_act)

tests/test_01_core_sample_locs.py

@@ -1,8 +1,10 @@
 import pytest
 from random import random
+from shapely.geometry import Point
 
 from pam.core import Population, Household, Person
 from pam.activity import Plan, Activity, Leg
+from pam.location import Location
 from .fixtures import *
 
 
@@ -57,3 +59,28 @@ def sample(self, location_idx, activity):
     population.sample_locs(FakeSampler())
     SmithHousehold[2].plan[2].location == SmithHousehold[2].plan[8].location
     SmithHousehold[2].plan[2].location == SmithHousehold[4].plan[2].location
+
+
+def test_retain_already_existing_locs(SmithHousehold):
+    """ The sampler does  """
+    population = Population()
+    population.add(SmithHousehold)
+    existing_location = Location(area='w', loc=Point(1, 1))
+
+    class FakeSampler:
+        def sample(self, location_idx, activity):
+            return random()
+
+    # keeps existing location, otherwise it samples
+    SmithHousehold[2].plan[2].location = existing_location
+    population.sample_locs(FakeSampler(), location_override=False)
+    for pid, person in SmithHousehold:
+        for i, act in enumerate(person.activities):
+            if pid==2 and i==1:
+                assert act.location == existing_location
+            else:
+                assert isinstance(act.location.loc, float)
+
+    # default behaviour is to override
+    population.sample_locs(FakeSampler())
+    assert SmithHousehold[2].plan[2].location != existing_location