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

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

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

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

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)