refactor facility index helper class

src/main/java/org/matsim/prepare/population/FacilityIndex.java

@@ -0,0 +1,57 @@
+package org.matsim.prepare.population;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.locationtech.jts.index.strtree.STRtree;
+import org.matsim.api.core.v01.Id;
+import org.matsim.core.utils.geometry.geotools.MGC;
+import org.matsim.facilities.ActivityFacilities;
+import org.matsim.facilities.ActivityFacility;
+import org.matsim.facilities.FacilitiesUtils;
+import org.matsim.facilities.MatsimFacilitiesReader;
+import org.matsim.run.RunOpenBerlinScenario;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Spatial index for facilities.
+ */
+final class FacilityIndex {
+
+	private static final Logger log = LogManager.getLogger(FacilityIndex.class);
+
+	final ActivityFacilities all = FacilitiesUtils.createActivityFacilities();
+
+	/**
+	 * Maps activity type to spatial index.
+	 */
+	final Map<String, STRtree> index = new HashMap<>();
+
+	public FacilityIndex(String facilityPath) {
+
+		new MatsimFacilitiesReader(RunOpenBerlinScenario.CRS, RunOpenBerlinScenario.CRS, all)
+			.readFile(facilityPath);
+
+		Set<String> activities = all.getFacilities().values().stream()
+			.flatMap(a -> a.getActivityOptions().keySet().stream())
+			.collect(Collectors.toSet());
+
+		log.info("Found activity types: {}", activities);
+
+		for (String act : activities) {
+
+			NavigableMap<Id<ActivityFacility>, ActivityFacility> afs = all.getFacilitiesForActivityType(act);
+			for (ActivityFacility af : afs.values()) {
+				STRtree index = this.index.computeIfAbsent(act, k -> new STRtree());
+				index.insert(MGC.coord2Point(af.getCoord()).getEnvelopeInternal(), af);
+			}
+		}
+
+		// Build all trees
+		index.values().forEach(STRtree::build);
+	}
+}

src/main/java/org/matsim/prepare/population/InitLocationChoice.java

@@ -82,17 +82,14 @@ public class InitLocationChoice implements MATSimAppCommand, PersonAlgorithm {
 	@CommandLine.Mixin
 	private ShpOptions shp;
 
-
-	private Map<String, STRtree> trees;
+	private FacilityIndex facilities;
 
 	private Long2ObjectMap<SimpleFeature> zones;
 
 	private CommuterAssignment commuter;
 
 	private Network network;
 
-	private ActivityFacilities facilities = FacilitiesUtils.createActivityFacilities();
-
 	private ThreadLocal<Context> ctxs;
 
 	private AtomicLong total = new AtomicLong();
@@ -131,28 +128,6 @@ public Integer call() throws Exception {
 		log.info("Read {} zones", zones.size());
 
 
-		new MatsimFacilitiesReader(RunOpenBerlinScenario.CRS, RunOpenBerlinScenario.CRS, facilities)
-			.readFile(facilityPath.toString());
-
-		Set<String> activities = facilities.getFacilities().values().stream()
-			.flatMap(a -> a.getActivityOptions().keySet().stream())
-			.collect(Collectors.toSet());
-
-		log.info("Found activity types: {}", activities);
-
-		trees = new HashMap<>();
-		for (String act : activities) {
-
-			NavigableMap<Id<ActivityFacility>, ActivityFacility> afs = facilities.getFacilitiesForActivityType(act);
-			for (ActivityFacility af : afs.values()) {
-				STRtree index = trees.computeIfAbsent(act, k -> new STRtree());
-				index.insert(MGC.coord2Point(af.getCoord()).getEnvelopeInternal(), af);
-			}
-		}
-
-		// Build all trees
-		trees.values().forEach(STRtree::build);
-
 		log.info("Using input file: {}", input);
 
 		List<Population> populations = new ArrayList<>();
@@ -241,11 +216,11 @@ public void run(Person person) {
 						location = sampleCommute(ctx, dist, lastCoord, (long) person.getAttributes().getAttribute(Attributes.ARS));
 					}
 
-					if (location == null && trees.containsKey(type)) {
+					if (location == null && facilities.index.containsKey(type)) {
 						// Needed for lambda
 						final Coord refCoord = lastCoord;
 
-						List<ActivityFacility> query = trees.get(type).query(MGC.coord2Point(lastCoord).buffer(dist * 1.2).getEnvelopeInternal());
+						List<ActivityFacility> query = facilities.index.get(type).query(MGC.coord2Point(lastCoord).buffer(dist * 1.2).getEnvelopeInternal());
 
 						// Distance should be within the bounds
 						List<ActivityFacility> res = query.stream().filter(f -> checkDistanceBound(dist, refCoord, f.getCoord(), 1)).toList();
@@ -271,7 +246,7 @@ public void run(Person person) {
 						lastCoord = c;
 
 						// An activity with type could not be put into correct facility.
-						if (trees.containsKey(type)) {
+						if (facilities.index.containsKey(type)) {
 							warning.incrementAndGet();
 						}
 
@@ -287,7 +262,7 @@ public void run(Person person) {
 				if (act.getCoord() != null)
 					lastCoord = act.getCoord();
 				else if (act.getFacilityId() != null)
-					lastCoord = facilities.getFacilities().get(act.getFacilityId()).getCoord();
+					lastCoord = facilities.all.getFacilities().get(act.getFacilityId()).getCoord();
 
 			}
 		}
@@ -298,7 +273,7 @@ else if (act.getFacilityId() != null)
 	 */
 	private ActivityFacility sampleCommute(Context ctx, double dist, Coord refCoord, long ars) {
 
-		STRtree index = trees.get("work");
+		STRtree index = facilities.index.get("work");
 
 		ActivityFacility workPlace = null;