Skip to content

Commit 9d597af

Browse files
samakinensamakinen
committed
Refactor geodata_helpers: streamline function definitions and imports
1 parent 0f20db5 commit 9d597af

File tree

1 file changed

+92
-98
lines changed

1 file changed

+92
-98
lines changed

‎Scripts/utils/geodata_helpers.py‎

Lines changed: 92 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -1,114 +1,108 @@
1-
from typing import Any, Dict, List, TYPE_CHECKING
2-
import pandas as pd
3-
from utils import log
1+
from typing import TYPE_CHECKING
42
if TYPE_CHECKING:
53
from assignment.emme_bindings.mock_project import Network
64
from assignment.emme_bindings.mock_project import Scenario
7-
try:
8-
import geopandas as gpd
9-
from shapely import LineString, Point
5+
import pandas as pd
6+
import geopandas as gpd
7+
from shapely import LineString, Point
108

11-
def get_links(network: 'Network', scen: 'Scenario') -> gpd.GeoDataFrame:
12-
"""
13-
Extracts link data from a network and scenario and returns it as a GeoDataFrame.
9+
def get_links(network: 'Network', scen: 'Scenario') -> gpd.GeoDataFrame:
10+
"""
11+
Extracts link data from a network and scenario and returns it as a GeoDataFrame.
1412
15-
Args:
16-
network (Network): The network object containing link data.
17-
scen (Scenario): The scenario object containing attributes and ID.
13+
Args:
14+
network (Network): The network object containing link data.
15+
scen (Scenario): The scenario object containing attributes and ID.
1816
19-
Returns:
20-
gpd.GeoDataFrame: A GeoDataFrame containing link data with the following columns:
21-
- 'scenario': The scenario ID.
22-
- 'id': The link ID.
23-
- 'geom': The geometry of the link as a LineString.
24-
- 'inode': The ID of the starting node of the link.
25-
- 'jnode': The ID of the ending node of the link.
26-
- Additional columns for each attribute in the scenario's "LINK" attributes.
27-
"""
28-
attributes = scen.attributes("LINK")
29-
obj = network.links()
30-
values = [dict(zip(
31-
['scenario', 'id', 'geom', 'inode', 'jnode'] + attributes,
32-
[int(scen.id), o.id, LineString(o.shape), o.i_node.id, o.j_node.id] + [o[a] for a in attributes]))
33-
for o in obj]
34-
return gpd.GeoDataFrame(values, geometry='geom').set_crs('epsg:3879')
17+
Returns:
18+
gpd.GeoDataFrame: A GeoDataFrame containing link data with the following columns:
19+
- 'scenario': The scenario ID.
20+
- 'id': The link ID.
21+
- 'geom': The geometry of the link as a LineString.
22+
- 'inode': The ID of the starting node of the link.
23+
- 'jnode': The ID of the ending node of the link.
24+
- Additional columns for each attribute in the scenario's "LINK" attributes.
25+
"""
26+
attributes = scen.attributes("LINK")
27+
obj = network.links()
28+
values = [dict(zip(
29+
['scenario', 'id', 'geom', 'inode', 'jnode'] + attributes,
30+
[int(scen.id), o.id, LineString(o.shape), o.i_node.id, o.j_node.id] + [o[a] for a in attributes]))
31+
for o in obj]
32+
return gpd.GeoDataFrame(values, geometry='geom').set_crs('epsg:3879')
3533

36-
def get_nodes(network: 'Network', scen: 'Scenario') -> gpd.GeoDataFrame:
37-
"""
38-
Extracts node data from a network and scenario, and returns it as a GeoDataFrame.
34+
def get_nodes(network: 'Network', scen: 'Scenario') -> gpd.GeoDataFrame:
35+
"""
36+
Extracts node data from a network and scenario, and returns it as a GeoDataFrame.
3937
40-
Args:
41-
network (Network): The network object containing node information.
42-
scen (Scenario): The scenario object containing attributes and scenario ID.
38+
Args:
39+
network (Network): The network object containing node information.
40+
scen (Scenario): The scenario object containing attributes and scenario ID.
4341
44-
Returns:
45-
gpd.GeoDataFrame: A GeoDataFrame containing node data with the following columns:
46-
- 'scenario': The scenario ID.
47-
- 'id': The node ID.
48-
- 'geom': The geometry of the node as a Point.
49-
- Additional columns for each attribute in the scenario's "NODE" attributes.
50-
"""
51-
attributes = scen.attributes("NODE")
52-
obj = network.nodes()
53-
values = [dict(zip(
54-
['scenario', 'id', 'geom'] + attributes,
55-
[int(scen.id), o.id, Point(o.x, o.y)] + [o[a] for a in attributes]))
56-
for o in obj]
57-
return gpd.GeoDataFrame(values, geometry='geom').set_crs('epsg:3879')
42+
Returns:
43+
gpd.GeoDataFrame: A GeoDataFrame containing node data with the following columns:
44+
- 'scenario': The scenario ID.
45+
- 'id': The node ID.
46+
- 'geom': The geometry of the node as a Point.
47+
- Additional columns for each attribute in the scenario's "NODE" attributes.
48+
"""
49+
attributes = scen.attributes("NODE")
50+
obj = network.nodes()
51+
values = [dict(zip(
52+
['scenario', 'id', 'geom'] + attributes,
53+
[int(scen.id), o.id, Point(o.x, o.y)] + [o[a] for a in attributes]))
54+
for o in obj]
55+
return gpd.GeoDataFrame(values, geometry='geom').set_crs('epsg:3879')
5856

59-
def get_transit_lines(network: 'Network', scen: 'Scenario') -> pd.DataFrame:
60-
"""
61-
Retrieve transit line data from a given network and scenario.
57+
def get_transit_lines(network: 'Network', scen: 'Scenario') -> pd.DataFrame:
58+
"""
59+
Retrieve transit line data from a given network and scenario.
6260
63-
This function extracts transit line information from the provided network
64-
and scenario, including attributes such as scenario ID, transit line ID,
65-
mode, vehicle, and any additional attributes defined for transit lines
66-
in the scenario.
61+
This function extracts transit line information from the provided network
62+
and scenario, including attributes such as scenario ID, transit line ID,
63+
mode, vehicle, and any additional attributes defined for transit lines
64+
in the scenario.
6765
68-
Args:
69-
network (Network): The network object containing transit line data.
70-
scen (Scenario): The scenario object containing attributes and context.
66+
Args:
67+
network (Network): The network object containing transit line data.
68+
scen (Scenario): The scenario object containing attributes and context.
7169
72-
Returns:
73-
pd.DataFrame: A DataFrame containing transit line data with the following columns:
74-
- 'scenario': The ID of the scenario.
75-
- 'id': The ID of the transit line.
76-
- 'mode': The mode of the transit line.
77-
- 'vehicle': The vehicle type of the transit line.
78-
- Additional columns for each attribute defined in the scenario for transit lines.
79-
"""
80-
attributes = scen.attributes("TRANSIT_LINE")
81-
obj = network.transit_lines()
82-
values = [dict(zip(
83-
['scenario', 'id', 'mode', 'vehicle'] + attributes,
84-
[int(scen.id), o.id, o.mode.id, o.vehicle.id] + [o[a] for a in attributes]))
85-
for o in obj]
86-
return pd.DataFrame(values)
70+
Returns:
71+
pd.DataFrame: A DataFrame containing transit line data with the following columns:
72+
- 'scenario': The ID of the scenario.
73+
- 'id': The ID of the transit line.
74+
- 'mode': The mode of the transit line.
75+
- 'vehicle': The vehicle type of the transit line.
76+
- Additional columns for each attribute defined in the scenario for transit lines.
77+
"""
78+
attributes = scen.attributes("TRANSIT_LINE")
79+
obj = network.transit_lines()
80+
values = [dict(zip(
81+
['scenario', 'id', 'mode', 'vehicle'] + attributes,
82+
[int(scen.id), o.id, o.mode.id, o.vehicle.id] + [o[a] for a in attributes]))
83+
for o in obj]
84+
return pd.DataFrame(values)
8785

88-
def get_transit_segments(network: 'Network', scen: 'Scenario') -> pd.DataFrame:
89-
"""
90-
Extracts transit segment data from a given network and scenario and returns it as a pandas DataFrame.
86+
def get_transit_segments(network: 'Network', scen: 'Scenario') -> pd.DataFrame:
87+
"""
88+
Extracts transit segment data from a given network and scenario and returns it as a pandas DataFrame.
9189
92-
Args:
93-
network (Network): The transit network object containing transit lines and segments.
94-
scen (Scenario): The scenario object containing attributes and scenario-specific data.
90+
Args:
91+
network (Network): The transit network object containing transit lines and segments.
92+
scen (Scenario): The scenario object containing attributes and scenario-specific data.
9593
96-
Returns:
97-
pd.DataFrame: A DataFrame containing transit segment data with the following columns:
98-
- 'scenario': The ID of the scenario.
99-
- 'id': The ID of the transit segment.
100-
- 'line_id': The ID of the transit line to which the segment belongs.
101-
- 'link_id': The ID of the link associated with the segment.
102-
- Additional columns for each attribute defined in the scenario's "TRANSIT_SEGMENT" attributes.
103-
"""
104-
attributes = scen.attributes("TRANSIT_SEGMENT")
105-
tls = network.transit_lines()
106-
values: List[Dict[str, Any]] = []
107-
for tl in tls:
108-
for ts in tl.segments():
109-
values.append(dict(zip(
110-
['scenario', 'id', 'line_id', 'link_id'] + attributes,
111-
[int(scen.id), ts.id, tl.id, ts.link.id] + [ts[a] for a in attributes])))
112-
return pd.DataFrame(values)
113-
except ImportError:
114-
log.warn("geopandas is not installed. Please install it to use this feature to export GPKG data.")
94+
Returns:
95+
pd.DataFrame: A DataFrame containing transit segment data with the following columns:
96+
- 'scenario': The ID of the scenario.
97+
- 'id': The ID of the transit segment.
98+
- 'line_id': The ID of the transit line to which the segment belongs.
99+
- 'link_id': The ID of the link associated with the segment.
100+
- Additional columns for each attribute defined in the scenario's "TRANSIT_SEGMENT" attributes.
101+
"""
102+
attributes = scen.attributes("TRANSIT_SEGMENT")
103+
tls = network.transit_lines()
104+
values = [dict(zip(
105+
['scenario', 'id', 'line_id', 'link_id'] + attributes,
106+
[int(scen.id), ts.id, tl.id, ts.link.id] + [ts[a] for a in attributes]))
107+
for tl in tls for ts in tl.segments()]
108+
return pd.DataFrame(values)

0 commit comments

Comments
 (0)