update ref data script

src/main/python/extract_ref_data.py

@@ -1,16 +1,20 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+import numpy as np
 import os
-
-from matsim.scenariogen.data import TripMode, run_create_ref_data
+from matsim.scenariogen.data import EconomicStatus, TripMode, preparation, run_create_ref_data
 
 
 def person_filter(df):
     """ Default person filter for reference data. """
     df = df[df.reporting_day <= 4]
     df = df[df.location == "Berlin"]
 
+    df["age_group"] = preparation.cut(df.age, [0, 18, 66, np.inf])
+
+    preparation.fill(df, "economic_status", EconomicStatus.UNKNOWN)
+
     return df
 
 
@@ -27,8 +31,13 @@ def trip_filter(df):
 if __name__ == "__main__":
     d = os.path.expanduser("~/Development/matsim-scenarios/shared-svn/projects/matsim-berlin/data/SrV/")
 
-    person, trips, share = run_create_ref_data.create(d + "Berlin+Umland",
-                                                      person_filter, trip_filter,
-                                                      run_create_ref_data.InvalidHandling.REMOVE_TRIPS)
+    result = run_create_ref_data.create(
+        d + "Berlin+Umland",
+        person_filter, trip_filter,
+        run_create_ref_data.InvalidHandling.REMOVE_TRIPS,
+        ref_groups=["age_group", "economic_status"]
+    )
+
+    print(result.share)
 
-    print(share)
+    print(result.groups)