Skip to content

Utils

This module provides utility functions for the project.

Functions

get_num_scenarios

Get total number of scenarios from data directory.

Reads from n_scenarios.txt metadata file in the data directory.

Parameters:

Name Type Description Default
data_dir str

Directory containing parquet files and n_scenarios.txt

 required

Returns:

Type Description
int

Total number of scenarios

Raises:

Type Description
ValueError

If n_scenarios.txt metadata file not found
Source code in gridfm_datakit/utils/utils.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def get_num_scenarios(data_dir: str) -> int:
    """Get total number of scenarios from data directory.

    Reads from n_scenarios.txt metadata file in the data directory.

    Args:
        data_dir: Directory containing parquet files and n_scenarios.txt

    Returns:
        Total number of scenarios

    Raises:
        ValueError: If n_scenarios.txt metadata file not found
    """
    n_scenarios_file = os.path.join(data_dir, "n_scenarios.txt")
    if os.path.exists(n_scenarios_file):
        with open(n_scenarios_file, "r") as f:
            return int(f.read().strip())

    else:
        print(
            f"No n_scenarios metadata file found in {data_dir}, using bus_data.parquet to get total number of scenarios",
        )
        return int(
            pd.read_parquet(
                os.path.join(data_dir, "bus_data.parquet"),
                engine="pyarrow",
            )["scenario"].max()
            + 1,
        )

read_partitions

Read sampled partition folders in parallel and concatenate them.

Source code in gridfm_datakit/utils/utils.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def read_partitions(
    base_path: str,
    sampled: list,
    max_workers: int = None,
) -> pd.DataFrame:
    """Read sampled partition folders in parallel and concatenate them."""
    if max_workers is None:
        from os import cpu_count

        max_workers = min(32, cpu_count())  # sensible default

    dfs = []

    # Submit all partition reads to the ThreadPool
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(
                pd.read_parquet,
                os.path.join(base_path, f"scenario_partition={k}"),
                engine="pyarrow",
            ): k
            for k in sampled
        }

        # Collect results as they complete with tqdm
        for future in tqdm(
            as_completed(futures),
            total=len(futures),
            desc=f"Reading {len(sampled)} partitions from {base_path}",
        ):
            df = future.result()
            dfs.append(df)

    # Concatenate all partitions
    return pd.concat(dfs, ignore_index=True)

write_ram_usage_distributed

Source code in gridfm_datakit/utils/utils.py 
79
80
81
82
83
84
85
86
87
88
89
90
def write_ram_usage_distributed(tqdm_log: TextIO) -> None:
    process = psutil.Process(os.getpid())  # Parent process
    mem_usage = process.memory_info().rss / 1024**2  # Parent memory in MB

    # Sum memory usage of all child processes
    for child in process.children(recursive=True):
        mem_usage += child.memory_info().rss / 1024**2

    _blocking_write(
        tqdm_log,
        f"Total RAM usage (Parent + Children): {mem_usage:.2f} MB\n",
    )

Classes

Tee

Source code in gridfm_datakit/utils/utils.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class Tee:
    def __init__(self, *streams):
        self.streams = streams

    def write(self, data):
        for s in self.streams:
            _blocking_write(s, data)
            _retry_on_eagain(s.flush)

    def flush(self):
        for s in self.streams:
            _retry_on_eagain(s.flush)