add method to sample with zone information

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();
+	}
 }

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()));
+	}
+
 }