re-worked the plan builder

src/main/java/org/matsim/prepare/RunOpenBerlinCalibration.java

@@ -44,7 +44,7 @@
 import org.matsim.core.scoring.SumScoringFunction;
 import org.matsim.core.scoring.functions.*;
 import org.matsim.core.utils.geometry.CoordUtils;
-import org.matsim.prepare.choices.AssignReferencePopulation;
+import org.matsim.prepare.population.AssignReferencePopulation;
 import org.matsim.prepare.choices.ComputePlanChoices;
 import org.matsim.prepare.choices.ComputeTripChoices;
 import org.matsim.prepare.counts.CreateCountsFromGeoPortalBerlin;

src/main/java/org/matsim/prepare/choices/PlanBuilder.java

@@ -34,7 +34,9 @@
 
 /**
  * Utility class to build plans from data.
+ * @deprecated use new plan builder in org.matsim.prepare.population
  */
+@Deprecated
 public class PlanBuilder {
 
 	private static final Logger log = LogManager.getLogger(PlanBuilder.class);

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

@@ -0,0 +1,138 @@
+package org.matsim.prepare.population;
+
+import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import me.tongfei.progressbar.ProgressBar;
+import org.apache.commons.csv.CSVRecord;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.matsim.api.core.v01.population.Person;
+import org.matsim.api.core.v01.population.Plan;
+import org.matsim.api.core.v01.population.Population;
+import org.matsim.application.MATSimAppCommand;
+import org.matsim.application.options.ShpOptions;
+import org.matsim.core.population.PopulationUtils;
+import org.matsim.core.router.TripStructureUtils;
+import picocli.CommandLine;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import java.util.SplittableRandom;
+
+@CommandLine.Command(
+	name = "assign-reference-population",
+	description = "Assigns persons from reference data to a population."
+)
+public class AssignReferencePopulation implements MATSimAppCommand {
+
+	private static final Logger log = LogManager.getLogger(AssignReferencePopulation.class);
+
+
+	@CommandLine.Option(names = "--population", description = "Input population path.", required = true)
+	private String populationPath;
+
+	@CommandLine.Option(names = "--persons", description = "Input persons from survey data, in matsim-python-tools format.", required = true)
+	private Path personsPath;
+
+	@CommandLine.Option(names = "--activities", description = "Path to activity table from survey data.", required = true)
+	private Path activityPath;
+
+	@CommandLine.Option(names = "--facilities", description = "Path to facilities file.", required = true)
+	private Path facilityPath;
+
+	@CommandLine.Option(names = "--network", description = "Path to network file", required = true)
+	private Path networkPath;
+
+	@CommandLine.Option(names = "--output", description = "Output population path.", required = true)
+	private Path output;
+
+	@CommandLine.Mixin
+	private ShpOptions shp;
+
+	private FacilityIndex facilities;
+
+	private PersonMatcher persons;
+
+	public static void main(String[] args) {
+		new AssignReferencePopulation().execute(args);
+	}
+
+	@Override
+	public Integer call() throws Exception {
+
+		if (!shp.isDefined()) {
+			log.error("No shapefile defined. Please specify a shapefile for the zones using the --shp option.");
+			return 2;
+		}
+
+		if (!Files.exists(activityPath)) {
+			log.error("Input activity file does not exist: {}", activityPath);
+			return 2;
+		}
+
+		Population population = PopulationUtils.readPopulation(populationPath);
+
+		SplittableRandom rnd = new SplittableRandom(0);
+		persons = new PersonMatcher("idx", personsPath);
+		facilities = new FacilityIndex(facilityPath.toString());
+
+		PlanBuilder planBuilder = new PlanBuilder(shp, facilities, activityPath);
+
+		Long2ObjectMap<List<Person>> homeIndex = planBuilder.createHomeIndex(population);
+
+		// Remove persons without legs, these can not be assigned
+		for (List<Person> list : homeIndex.values()) {
+			list.removeIf(p -> TripStructureUtils.getLegs(p.getSelectedPlan()).isEmpty());
+		}
+
+		RunActivitySampling sampling = new RunActivitySampling(persons, planBuilder.getActivities(), population.getFactory(), 1);
+
+		for (Map.Entry<String, CSVRecord> e : ProgressBar.wrap(persons, "Assigning reference population")) {
+
+			CSVRecord p = e.getValue();
+			if (!p.get("seq").equals("0"))
+				continue;
+
+			if (p.get("mobile_on_day").equals("false"))
+				continue;
+
+			long zone = planBuilder.findHomeZone(e.getKey());
+
+			// No home zone known
+			if (zone < 0)
+				continue;
+
+			List<Person> refPersons = homeIndex.get(zone);
+
+			if (refPersons == null)
+				continue;
+
+			Person person = persons.matchEntry(e.getValue(), refPersons, rnd);
+
+			// No persons matched
+			if (person == null)
+				continue;
+
+			// Create the base daily plan (without locations)
+			Plan plan = sampling.createPlan(Attributes.getHomeCoord(person), e.getKey());
+
+			boolean success = planBuilder.assignLocationsFromZones(e.getKey(), plan);
+
+			if (success) {
+				sampling.copyAttributes(p, person);
+				person.getAttributes().putAttribute(Attributes.REF_WEIGHT, p.get("p_weight"));
+				person.removePlan(person.getSelectedPlan());
+				person.addPlan(plan);
+				person.setSelectedPlan(plan);
+
+				// remove person that have been used as reference
+				refPersons.remove(person);
+			}
+		}
+
+		PopulationUtils.writePopulation(population, output.toString());
+
+		return 0;
+	}
+}

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

@@ -36,6 +36,8 @@ public final class Attributes {
 	public static final String RESTRICTED_MOBILITY = "restricted_mobility";
 	public static final String ECONOMIC_STATUS = "economic_status";
 	public static final String HOUSEHOLD_SIZE = "household_size";
+	public static final String REF_WEIGHT = "ref_weight";
+	public static final String REF_MODES = "ref_modes";
 	public static final String ATTRACTION_WORK = "attraction_work";
 	public static final String ATTRACTION_OTHER = "attraction_other";
 

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

@@ -5,6 +5,7 @@
 import org.apache.commons.csv.CSVRecord;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.jetbrains.annotations.NotNull;
 import org.matsim.api.core.v01.population.Person;
 import org.matsim.application.options.CsvOptions;
 import org.matsim.core.population.PersonUtils;
@@ -13,13 +14,14 @@
 import java.io.UncheckedIOException;
 import java.nio.file.Path;
 import java.util.*;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 /**
  * This class is used to read and match persons from the reference data in csv format.
  */
-public class PersonMatcher {
+public class PersonMatcher implements Iterable<Map.Entry<String, CSVRecord>> {
 
 	private static final Logger log = LogManager.getLogger(PersonMatcher.class);
 
@@ -60,6 +62,29 @@ public String matchPerson(Person person, SplittableRandom rnd) {
 		return subgroup.get(rnd.nextInt(subgroup.size()));
 	}
 
+	/**
+	 * Matches a person csv entry to one person in a list.
+	 * @return null if no match was found
+	 */
+	public Person matchEntry(CSVRecord p, List<Person> refPersons, SplittableRandom rnd) {
+
+		int regionType = Integer.parseInt(p.get("region_type"));
+		String gender = p.get("gender");
+		String employment = p.get("employment");
+		int age = Integer.parseInt(p.get("age"));
+
+		Set<Key> keys = createKey(gender, age, regionType, employment).collect(Collectors.toSet());
+
+		List<Person> matched = refPersons.stream()
+			.filter(person -> keys.contains(createKey(person)))
+			.toList();
+
+		if (matched.isEmpty())
+			return null;
+
+		return matched.get(rnd.nextInt(matched.size()));
+	}
+
 	/**
 	 * Return reference person with given index.
 	 */
@@ -135,6 +160,12 @@ private Key createKey(Person person) {
 		return new Key(gender, age, regionType, employed);
 	}
 
+	@NotNull
+	@Override
+	public Iterator<Map.Entry<String, CSVRecord>> iterator() {
+		return persons.entrySet().iterator();
+	}
+
 	/**
 	 * Key used to match persons.
 	 */