From 8e9c57f5da804282156a4d5d9091811be1e1c6da Mon Sep 17 00:00:00 2001
From: rakow <rakow@vsp.tu-berlin.de>
Date: Fri, 9 Aug 2024 22:52:29 +0200
Subject: [PATCH] add method to sample with zone information

---
 .../AttributedActivityFacility.java           | 17 ++++++
 .../prepare/population/FacilityIndex.java     | 60 ++++++++++++++++++-
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/matsim/prepare/facilities/AttributedActivityFacility.java b/src/main/java/org/matsim/prepare/facilities/AttributedActivityFacility.java
index dc76990c..238dfa69 100644
--- a/src/main/java/org/matsim/prepare/facilities/AttributedActivityFacility.java
+++ b/src/main/java/org/matsim/prepare/facilities/AttributedActivityFacility.java
@@ -23,10 +23,15 @@ public final class AttributedActivityFacility implements ActivityFacility {
 	private final double workAttraction;
 	private final double otherAttraction;
 
+	private final String location;
+	private final String zone;
+
 	public AttributedActivityFacility(ActivityFacility facility) {
 		this.facility = facility;
 		this.workAttraction = (double) facility.getAttributes().getAttribute(ATTRACTION_WORK);
 		this.otherAttraction = (double) facility.getAttributes().getAttribute(ATTRACTION_OTHER);
+		this.location = (String) facility.getAttributes().getAttribute("location");
+		this.zone = (String) facility.getAttributes().getAttribute("zone");
 	}
 
 	public double getWorkAttraction() {
@@ -76,4 +81,16 @@ public Map<String, Object> getCustomAttributes() {
 	public Attributes getAttributes() {
 		return facility.getAttributes();
 	}
+
+	public String getLocation() {
+		return location;
+	}
+
+	public String getZone() {
+		return zone;
+	}
+
+	public String toString() {
+		return facility.toString();
+	}
 }
diff --git a/src/main/java/org/matsim/prepare/population/FacilityIndex.java b/src/main/java/org/matsim/prepare/population/FacilityIndex.java
index bc4b6993..97cca6a2 100644
--- a/src/main/java/org/matsim/prepare/population/FacilityIndex.java
+++ b/src/main/java/org/matsim/prepare/population/FacilityIndex.java
@@ -22,19 +22,27 @@
 public final class FacilityIndex {
 
 	private static final Logger log = LogManager.getLogger(FacilityIndex.class);
-
-	final ActivityFacilities all = FacilitiesUtils.createActivityFacilities();
-
 	/**
 	 * Maps activity type to spatial index.
 	 */
 	public final Map<String, STRtree> index = new HashMap<>();
+	final ActivityFacilities all = FacilitiesUtils.createActivityFacilities();
 
 	public FacilityIndex(String facilityPath, String crs) {
+		this(facilityPath, f -> true, crs);
+	}
+
+	/**
+	 * Creates spatial index for facilities.
+	 *
+	 * @param f predicate to filter or transform facilities
+	 */
+	public FacilityIndex(String facilityPath, Predicate<ActivityFacility> f, String crs) {
 
 		new MatsimFacilitiesReader(crs, crs, all).readFile(facilityPath);
 
 		Set<String> activities = all.getFacilities().values().stream()
+			.filter(f)
 			.flatMap(a -> a.getActivityOptions().keySet().stream())
 			.collect(Collectors.toSet());
 
@@ -127,4 +135,50 @@ public static ActivityFacility sampleByWeightWithRejection(List<AttributedActivi
 		return null;
 	}
 
+	/**
+	 * Groups candidates first using classifier. Then does a weighted sample on the groups and selects a random facility.
+	 * The groups are typical zones which might have certain OD relations.
+	 */
+	public static ActivityFacility sampleWithGrouping(List<AttributedActivityFacility> candidates,
+											 Function<AttributedActivityFacility, String> classifier,
+											 Function<Map.Entry<String, List<AttributedActivityFacility>>, Double> groupWeight,
+											 SplittableRandom rnd) {
+
+		if (candidates.isEmpty())
+			return null;
+
+		Map<String, List<AttributedActivityFacility>> map = candidates.stream()
+			.filter(f -> f.getLocation() != null)
+			.collect(Collectors.groupingBy(classifier));
+
+		List<Map.Entry<String, List<AttributedActivityFacility>>> grouped = map.entrySet().stream().toList();
+
+		double totalWeight = 0.0;
+		double[] weights = new double[grouped.size()];
+
+		for (int i = 0; i < grouped.size(); ++i) {
+			double w = groupWeight.apply(grouped.get(i));
+			totalWeight += w;
+			weights[i] = totalWeight;
+		}
+
+		// No weights, sample uniformly
+		if (totalWeight == 0.0) {
+			List<AttributedActivityFacility> list = grouped.get(rnd.nextInt(grouped.size())).getValue();
+			return list.get(rnd.nextInt(list.size()));
+		}
+
+		double r = rnd.nextDouble(0.0, totalWeight);
+		int idx = Arrays.binarySearch(weights, r);
+		if (idx < 0) {
+			idx = -idx - 1;
+		}
+
+		// First sample a random group.
+		List<AttributedActivityFacility> list = grouped.get(idx).getValue();
+
+		// Sample random facility from the zone
+		return list.get(rnd.nextInt(list.size()));
+	}
+
 }