From ed2e526dcd171cd96c79e687b1ad99b47e8292b1 Mon Sep 17 00:00:00 2001
From: rakow <rakow@vsp.tu-berlin.de>
Date: Sun, 18 Feb 2024 12:00:51 +0100
Subject: [PATCH] add zone information to data model

---
 matsim/scenariogen/data/__init__.py    |  4 +++
 matsim/scenariogen/data/formats/srv.py | 38 ++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/matsim/scenariogen/data/__init__.py b/matsim/scenariogen/data/__init__.py
index 230d914..5f36c0f 100644
--- a/matsim/scenariogen/data/__init__.py
+++ b/matsim/scenariogen/data/__init__.py
@@ -288,6 +288,8 @@ class Household:
     type: HouseholdType
     region_type: int
     location: str
+    zone: str = None
+    """ A detailed zone, which can be more accurate than location. """
     income: float = None
     geom: object = None
 
@@ -328,6 +330,8 @@ class Trip:
     purpose: Purpose
     sd_group: SourceDestinationGroup
     valid: bool
+    from_zone: str = None
+    to_zone: str = None
 
 
 @dataclass
diff --git a/matsim/scenariogen/data/formats/srv.py b/matsim/scenariogen/data/formats/srv.py
index bd02d1a..79bb094 100644
--- a/matsim/scenariogen/data/formats/srv.py
+++ b/matsim/scenariogen/data/formats/srv.py
@@ -2,14 +2,15 @@
 
 import os
 
-import pandas as pd
 import numpy as np
+import pandas as pd
 
 from .. import *
 
 # Has households, persons and trips
 INPUT_FILES = 3
 
+
 def is_format(f: os.DirEntry):
     fp = f.name
     if not f.path.endswith(".csv"):
@@ -92,6 +93,7 @@ def convert(data: tuple, regio=None):
                 SrV2018.household_type(h.E_HHTYP),
                 SrV2018.region_type(h, regio, random_state),
                 h.ST_CODE_NAME,
+                zone=SrV2018.parse_zone(h),
                 income=SrV2018.income(h.V_EINK),
             )
         )
@@ -116,17 +118,33 @@ def convert(data: tuple, regio=None):
                 SrV2018.trip_purpose(t.V_ZWECK),
                 SrV2018.sd_group(int(t.E_QZG_17)),
                 # Trip is valid if length and duration are present
-                0 <= t.GIS_LAENGE and t.E_DAUER > 0
+                0 <= t.GIS_LAENGE and t.E_DAUER > 0,
+                from_zone=SrV2018.parse_zone(t, "V_START_"),
+                to_zone=SrV2018.parse_zone(t, "V_ZIEL_")
             )
         )
 
     return pd.DataFrame(hhs).set_index("hh_id"), ps, pd.DataFrame(ts).set_index("t_id")
 
+
 def pint(x):
     """ Convert to positive integer"""
     return max(0, int(x))
 
 
+def parse_int_str(x):
+    """ Return parsed int or string """
+    try:
+        i = int(x)
+        if i >= 0:
+            return str(i)
+        return None
+    except ValueError:
+        if not x:
+            return None
+        return x
+
+
 class SrV2018:
     """ Maps SrV data to standard format"""
 
@@ -420,3 +438,19 @@ def income(x):
             return 5600
 
         return -1
+
+    @staticmethod
+    def parse_zone(h, prefix=""):
+        ob = parse_int_str(getattr(h, prefix + "OBERBEZIRK"))
+
+        zone = None
+        if ob:
+            zone = ob
+            ub = parse_int_str(getattr(h, prefix + "UNTERBEZIRK"))
+            if ub:
+                zone += "-" + ub
+                tb = parse_int_str(getattr(h, prefix + "TEILBEZIRK"))
+                if tb:
+                    zone += "-" + tb
+
+        return zone