Utils Io

io

CLASS	DESCRIPTION
DataFormat	Supported data formats for loading/saving.

FUNCTION	DESCRIPTION
detect_format	Auto-detect data file format.
create_dataset	Create DatasetWeighted from various data sources.
load_data	Load data from files with automatic format detection.
skip_every_nth	Create a function that skips every nth row.
keep_every_nth	Create a function that keeps every nth row.
save_data	Save data in specified format.
save_object	Serialize an object using pickle.
load_object	Load serialized object from pickle file.
save_metrics	Save metrics dict to file.

Classes

DataFormat

Bases: Enum

Supported data formats for loading/saving.

Functions

detect_format(filename: str) -> DataFormat

Auto-detect data file format.

PARAMETER	DESCRIPTION
filename	File path to analyze. TYPE: str

RETURNS	DESCRIPTION
DataFormat	Detected format.

Source code in spectre/utils/io.py

def detect_format(filename: str) -> DataFormat:
    """
    Auto-detect data file format.

    Parameters
    ----------
    filename : str
        File path to analyze.

    Returns
    -------
    DataFormat
        Detected format.
    """
    path = pathlib.Path(filename)

    # Check file extension
    suffix = path.suffix.lower()
    if suffix == ".npy":
        return DataFormat.NUMPY
    elif suffix == ".csv":
        return DataFormat.CSV
    elif suffix in [".tsv", ".tab"]:
        return DataFormat.TSV

    # Check for PLUMED format by examining file content
    try:
        with open(filename, "r") as f:
            first_line = f.readline().strip()
            if first_line.startswith("#! FIELDS"):
                return DataFormat.PLUMED
    except (OSError, UnicodeDecodeError):
        pass

    # Default to plain text
    return DataFormat.TEXT

create_dataset(data_source: str | list[str] | pd.DataFrame | np.ndarray | torch.Tensor, target_cols: str | list[str] | None = None, weight_cols: str | list[str] | None = None, weight_reduce: Literal['prod', 'sum'] = 'prod', **load_kwargs) -> DatasetWeighted

Create DatasetWeighted from various data sources.

PARAMETER	DESCRIPTION
data_source	Data source - file path(s) or data array/DataFrame. TYPE: str | list[str] | DataFrame | ndarray | Tensor
target_cols	Target column names/indices, by default None. TYPE: str | list[str] | None DEFAULT: None
weight_cols	Weight column names/indices, by default None. TYPE: str | list[str] | None DEFAULT: None
weight_reduce	Weight reduction method, by default "prod". TYPE: Literal['prod', 'sum'] DEFAULT: 'prod'
**load_kwargs	Additional arguments passed to load_data(). Common options: - skip_rows: int | Callable for row skipping - format: DataFormat for explicit format specification - delimiter: str for custom column separators DEFAULT: {}

RETURNS	DESCRIPTION
DatasetWeighted	Created dataset with data, weights, and targets.

Examples:

>>> # From files
>>> dataset = create_dataset("data.txt", target_cols="target", weight_cols="weight")
>>>
>>> # From PLUMED files
>>> dataset = create_dataset(["file1.dat", "file2.dat"], format="plumed")
>>>
>>> # From DataFrame
>>> dataset = create_dataset(df, target_cols=["y"], weight_cols=["w"])

Source code in spectre/utils/io.py

def create_dataset(
    data_source: str | list[str] | pd.DataFrame | np.ndarray | torch.Tensor,
    target_cols: str | list[str] | None = None,
    weight_cols: str | list[str] | None = None,
    weight_reduce: Literal["prod", "sum"] = "prod",
    **load_kwargs,
) -> DatasetWeighted:
    """
    Create DatasetWeighted from various data sources.

    Parameters
    ----------
    data_source : str | list[str] | pd.DataFrame | np.ndarray | torch.Tensor
        Data source - file path(s) or data array/DataFrame.
    target_cols : str | list[str] | None, optional
        Target column names/indices, by default None.
    weight_cols : str | list[str] | None, optional
        Weight column names/indices, by default None.
    weight_reduce : Literal["prod", "sum"], optional
        Weight reduction method, by default "prod".
    **load_kwargs
        Additional arguments passed to load_data(). Common options:
        - skip_rows: int | Callable for row skipping
        - format: DataFormat for explicit format specification
        - delimiter: str for custom column separators

    Returns
    -------
    DatasetWeighted
        Created dataset with data, weights, and targets.

    Examples
    --------
    >>> # From files
    >>> dataset = create_dataset("data.txt", target_cols="target", weight_cols="weight")
    >>>
    >>> # From PLUMED files
    >>> dataset = create_dataset(["file1.dat", "file2.dat"], format="plumed")
    >>>
    >>> # From DataFrame
    >>> dataset = create_dataset(df, target_cols=["y"], weight_cols=["w"])
    """
    # Load data if needed
    if isinstance(data_source, (str, list)):
        df = load_data(data_source, **load_kwargs)
    elif isinstance(data_source, pd.DataFrame):
        df = data_source.copy()
    elif isinstance(data_source, (np.ndarray, torch.Tensor)):
        df = pd.DataFrame(data_source)
    else:
        raise TypeError(f"Unsupported data source type: {type(data_source)}.")

    # Extract targets
    target = None
    if target_cols is not None:
        if isinstance(target_cols, str):
            target_cols = [target_cols]

        target_data = df[target_cols].values
        target = torch.from_numpy(target_data.astype(np.float32))
        if target.shape[1] == 1:
            target = target.squeeze(1)

        df = df.drop(columns=target_cols)

    # Extract weights
    weights = None
    if weight_cols is not None:
        if isinstance(weight_cols, str):
            weight_cols = [weight_cols]

        weight_data = df[weight_cols].values
        weights = torch.from_numpy(weight_data.astype(np.float32))
        if weights.shape[1] == 1:
            weights = weights.squeeze(1)

        df = df.drop(columns=weight_cols)

    # Remaining columns are data
    data = torch.from_numpy(df.values.astype(np.float32))

    # Extract feature labels from remaining DataFrame columns
    feature_labels = list(df.columns)

    return DatasetWeighted(
        X=data,
        target=target,
        weights=weights,
        reduce=weight_reduce,
        labels=feature_labels if feature_labels else None,
    )

load_data(filename: str | list[str], format: DataFormat | str = DataFormat.AUTO, col_names: list[str] | None = None, skip_columns: int | list[int] | None = None, select_columns: list[int] | list[str] | str | None = None, skip_rows: int | Callable | None = None, delimiter: str | None = None, header: bool | None = None) -> pd.DataFrame

Load data from files with automatic format detection.

PARAMETER	DESCRIPTION
filename	Single file or list of files to load and concatenate. TYPE: str | list[str]
format	Data format, by default DataFormat.AUTO for auto-detection. TYPE: DataFormat | str DEFAULT: AUTO
col_names	Column names for the data, by default None. TYPE: list[str] | None DEFAULT: None
skip_columns	Column indices to skip, by default None. TYPE: int | list[int] | None DEFAULT: None
select_columns	Columns to extract. Can be: list[int]: Column indices to select list[str]: Column names to select (requires header or col_names) str: Glob pattern for column names (e.g., "d-*" for columns starting with "d-", "*_mean" for columns ending with "_mean") By default None (keep all columns). TYPE: list[int] | list[str] | str | None DEFAULT: None
skip_rows	Row skipping specification, by default None. Can be: - int: Skip first N rows - Callable: Function that takes row index and returns True to skip Examples: lambda i: i % 2 == 0 (skip even rows), lambda i: i % 10 != 0 (keep every 10th row) TYPE: int | Callable | None DEFAULT: None
delimiter	Column delimiter (auto-detected if None), by default None. TYPE: str | None DEFAULT: None
header	Whether first row is header (auto-detected if None), by default None. TYPE: bool | None DEFAULT: None

RETURNS	DESCRIPTION
DataFrame	Loaded data as DataFrame.

RAISES	DESCRIPTION
ValueError	If files cannot be loaded or have incompatible formats.

Examples:

>>> # Load specific columns by index
>>> df = load_data("data.txt", select_columns=[0, 2, 5])
>>>
>>> # Load columns by name
>>> df = load_data("data.csv", select_columns=["time", "temperature", "pressure"])
>>>
>>> # Load columns matching pattern
>>> df = load_data("data.txt", select_columns="d-*")  # columns starting with "d-"

Source code in spectre/utils/io.py

def load_data(
    filename: str | list[str],
    format: DataFormat | str = DataFormat.AUTO,
    col_names: list[str] | None = None,
    skip_columns: int | list[int] | None = None,
    select_columns: list[int] | list[str] | str | None = None,
    skip_rows: int | Callable | None = None,
    delimiter: str | None = None,
    header: bool | None = None,
) -> pd.DataFrame:
    """
    Load data from files with automatic format detection.

    Parameters
    ----------
    filename : str | list[str]
        Single file or list of files to load and concatenate.
    format : DataFormat | str, optional
        Data format, by default DataFormat.AUTO for auto-detection.
    col_names : list[str] | None, optional
        Column names for the data, by default None.
    skip_columns : int | list[int] | None, optional
        Column indices to skip, by default None.
    select_columns : list[int] | list[str] | str | None, optional
        Columns to extract. Can be:

        - list[int]: Column indices to select
        - list[str]: Column names to select (requires header or `col_names`)
        - str: Glob pattern for column names (e.g., "d-`*`" for columns starting
          with "d-", "`*`_mean" for columns ending with "_mean")

        By default None (keep all columns).
    skip_rows : int | Callable | None, optional
        Row skipping specification, by default None. Can be:
        - int: Skip first N rows
        - Callable: Function that takes row index and returns True to skip
        Examples: lambda i: i % 2 == 0 (skip even rows), lambda i: i % 10 != 0
        (keep every 10th row)
    delimiter : str | None, optional
        Column delimiter (auto-detected if None), by default None.
    header : bool | None, optional
        Whether first row is header (auto-detected if None), by default None.

    Returns
    -------
    pd.DataFrame
        Loaded data as DataFrame.

    Raises
    ------
    ValueError
        If files cannot be loaded or have incompatible formats.

    Examples
    --------
    >>> # Load specific columns by index
    >>> df = load_data("data.txt", select_columns=[0, 2, 5])
    >>>
    >>> # Load columns by name
    >>> df = load_data("data.csv", select_columns=["time", "temperature", "pressure"])
    >>>
    >>> # Load columns matching pattern
    >>> df = load_data("data.txt", select_columns="d-*")  # columns starting with "d-"
    """
    # Normalize inputs
    if isinstance(filename, str):
        filename = [filename]

    if isinstance(format, str):
        format = DataFormat(format)

    if isinstance(skip_columns, int):
        skip_columns = [skip_columns]

    # Validate mutually exclusive options
    if skip_columns and select_columns:
        raise ValueError(
            "Cannot specify both `skip_columns` and `select_columns`. "
            "Use one or the other."
        )

    # Load and concatenate files
    dataframes = []
    for file in filename:
        df = _load_single_file(file, format, delimiter, header, skip_rows)
        dataframes.append(df)

    try:
        combined_df = pd.concat(dataframes, axis=0, ignore_index=True)
    except ValueError as e:
        raise ValueError(
            "Cannot concatenate files. Check if they have compatible shapes."
        ) from e

    # Apply column names if provided (before selection)
    if col_names:
        if len(col_names) != len(combined_df.columns):
            raise ValueError(
                f"Column names length ({len(col_names)}) does not match "
                f"data columns ({len(combined_df.columns)})."
            )
        combined_df.columns = col_names

    # Apply column selection
    if select_columns is not None:
        if isinstance(select_columns, list):
            if all(isinstance(col, int) for col in select_columns):
                # Select by indices
                combined_df = combined_df.iloc[:, select_columns]
            elif all(isinstance(col, str) for col in select_columns):
                # Select by names
                combined_df = combined_df[select_columns]
            else:
                raise TypeError(
                    "`select_columns` list must contain all ints or all strings."
                )
        elif isinstance(select_columns, str):
            # Use pandas filter with regex - convert glob pattern to regex
            regex_pattern = select_columns.replace("*", ".*").replace("?", ".")
            regex_pattern = "^" + regex_pattern + "$"
            combined_df = combined_df.filter(regex=regex_pattern)
            if combined_df.empty or len(combined_df.columns) == 0:
                raise ValueError(f"No columns matched the pattern '{select_columns}'.")
        else:
            raise TypeError("`select_columns` must be list[int], list[str], or str.")

    # Apply column skipping (legacy support)
    if skip_columns:
        combined_df = combined_df.drop(combined_df.columns[skip_columns], axis=1)

    return combined_df

skip_every_nth(n: int) -> Callable

Create a function that skips every nth row.

PARAMETER	DESCRIPTION
n	Skip every nth row (keep rows where index % n != 0). TYPE: int

RETURNS	DESCRIPTION
Callable	Function that returns True for rows to skip.

Examples:

>>> # Keep every 10th row (skip 9 out of 10)
>>> skip_func = skip_every_nth(10)
>>> data = load_data("large_file.txt", skip_rows=skip_func)

Source code in spectre/utils/io.py

def skip_every_nth(n: int) -> Callable:
    """Create a function that skips every nth row.

    Parameters
    ----------
    n : int
        Skip every nth row (keep rows where index % n != 0).

    Returns
    -------
    Callable
        Function that returns True for rows to skip.

    Examples
    --------
    >>> # Keep every 10th row (skip 9 out of 10)
    >>> skip_func = skip_every_nth(10)
    >>> data = load_data("large_file.txt", skip_rows=skip_func)
    """
    return lambda i: i % n != 0

keep_every_nth(n: int) -> Callable

Create a function that keeps every nth row.

PARAMETER	DESCRIPTION
n	Keep every nth row (skip rows where index % n != 0). TYPE: int

RETURNS	DESCRIPTION
Callable	Function that returns True for rows to skip.

Examples:

>>> # Keep every 5th row
>>> skip_func = keep_every_nth(5)
>>> data = load_data("file.txt", skip_rows=skip_func)

Source code in spectre/utils/io.py

def keep_every_nth(n: int) -> Callable:
    """Create a function that keeps every nth row.

    Parameters
    ----------
    n : int
        Keep every nth row (skip rows where index % n != 0).

    Returns
    -------
    Callable
        Function that returns True for rows to skip.

    Examples
    --------
    >>> # Keep every 5th row
    >>> skip_func = keep_every_nth(5)
    >>> data = load_data("file.txt", skip_rows=skip_func)
    """
    return lambda i: i % n != 0 and i != 0  # Keep row 0 and every nth row

save_data(data: np.ndarray | torch.Tensor | pd.DataFrame, filename: str, format: DataFormat | str = DataFormat.AUTO, col_names: list[str] | None = None, fmt: str = '%.6f', delimiter: str | None = None, header: bool = True) -> None

Save data in specified format.

PARAMETER	DESCRIPTION
data	Data to save. TYPE: ndarray | Tensor | DataFrame
filename	Output file path. TYPE: str
format	Output format, by default DataFormat.AUTO (inferred from filename). TYPE: DataFormat | str DEFAULT: AUTO
col_names	Column names for output, by default None. TYPE: list[str] | None DEFAULT: None
fmt	Number format string, by default "%.6f". TYPE: str DEFAULT: '%.6f'
delimiter	Column delimiter, by default None (format-specific default). TYPE: str | None DEFAULT: None
header	Include header in output, by default True. TYPE: bool DEFAULT: True

RAISES	DESCRIPTION
ValueError	If data format is invalid or incompatible.
OSError	If file cannot be written.

Source code in spectre/utils/io.py

def save_data(
    data: np.ndarray | torch.Tensor | pd.DataFrame,
    filename: str,
    format: DataFormat | str = DataFormat.AUTO,
    col_names: list[str] | None = None,
    fmt: str = "%.6f",
    delimiter: str | None = None,
    header: bool = True,
) -> None:
    """
    Save data in specified format.

    Parameters
    ----------
    data : np.ndarray | torch.Tensor | pd.DataFrame
        Data to save.
    filename : str
        Output file path.
    format : DataFormat | str, optional
        Output format, by default DataFormat.AUTO (inferred from filename).
    col_names : list[str] | None, optional
        Column names for output, by default None.
    fmt : str, optional
        Number format string, by default "%.6f".
    delimiter : str | None, optional
        Column delimiter, by default None (format-specific default).
    header : bool, optional
        Include header in output, by default True.

    Raises
    ------
    ValueError
        If data format is invalid or incompatible.
    OSError
        If file cannot be written.
    """
    if isinstance(format, str):
        format = DataFormat(format)

    if format == DataFormat.AUTO:
        format = _infer_save_format(filename)

    # Convert data to numpy if needed
    if torch.is_tensor(data):
        data = data.detach().cpu().numpy()
    elif isinstance(data, pd.DataFrame):
        if col_names is None:
            col_names = list(data.columns)
        data = data.values

    if data.ndim != 2:
        raise ValueError(f"Expected 2D data, got {data.ndim}D.")

    # Generate column names if needed
    if col_names is None and header:
        if format == DataFormat.PLUMED:
            col_names = [f"v{i}" for i in range(1, data.shape[1] + 1)]
        else:
            col_names = [f"col_{i}" for i in range(data.shape[1])]

    try:
        if format == DataFormat.NUMPY:
            np.save(filename, data)

        elif format == DataFormat.PLUMED:
            _save_plumed_format(filename, data, col_names, fmt)

        elif format == DataFormat.CSV:
            _save_delimited_format(filename, data, col_names, ",", fmt, header)

        elif format == DataFormat.TSV:
            _save_delimited_format(filename, data, col_names, "\t", fmt, header)

        else:  # DataFormat.TEXT
            delim = delimiter or " "
            _save_delimited_format(filename, data, col_names, delim, fmt, header)

    except (OSError, PermissionError) as e:
        raise OSError(f"Failed to save data to '{filename}': {e}") from e

save_object(obj: Any, filename: str) -> None

Serialize an object using pickle.

Source code in spectre/utils/io.py

def save_object(obj: Any, filename: str) -> None:
    """Serialize an object using pickle."""
    try:
        with open(filename, "wb") as output:
            pickle.dump(obj, output, pickle.HIGHEST_PROTOCOL)
    except (OSError, PermissionError) as e:
        raise OSError(f"Failed to save object to '{filename}': {e}.") from e

load_object(filename: str) -> Any

Load serialized object from pickle file.

Source code in spectre/utils/io.py

def load_object(filename: str) -> Any:
    """Load serialized object from pickle file."""
    try:
        with open(filename, "rb") as input_file:
            obj = pickle.load(input_file)
        return obj
    except (OSError, FileNotFoundError, pickle.PickleError) as e:
        raise OSError(f"Failed to load object from '{filename}': {e}.") from e

save_metrics(metrics: dict, filename: str | Path, format: str | None = None) -> None

Save metrics dict to file.

PARAMETER	DESCRIPTION
metrics	Metrics to save. TYPE: dict
filename	Output file path. TYPE: str | Path
format	Save format ("pickle" or "json"). If None, uses save_format attribute, or auto-detects from file extension. TYPE: str | None DEFAULT: None

Notes

Format auto-detection rules: - .json extension → JSON format - .pkl, .pickle extension → pickle format - Otherwise → uses save_format attribute (default: pickle)

Source code in spectre/utils/io.py

def save_metrics(
    metrics: dict, filename: str | Path, format: str | None = None
) -> None:
    """
    Save metrics dict to file.

    Parameters
    ----------
    metrics : dict
        Metrics to save.
    filename : str | Path
        Output file path.
    format : str | None, optional
        Save format ("pickle" or "json"). If None, uses `save_format`
        attribute, or auto-detects from file extension.

    Notes
    -----
    Format auto-detection rules:
    - `.json` extension → JSON format
    - `.pkl`, `.pickle` extension → pickle format
    - Otherwise → uses `save_format` attribute (default: pickle)
    """
    filename = Path(filename)

    # Determine format: explicit > file extension > default
    if format is None:
        ext = filename.suffix.lower()
        if ext == ".json":
            format = "json"
        elif ext in [".pkl", ".pickle"]:
            format = "pickle"

    # Ensure parent directory exists
    filename.parent.mkdir(parents=True, exist_ok=True)

    if format == "pickle":
        with open(filename, "wb") as f:
            pickle.dump(metrics, f, pickle.HIGHEST_PROTOCOL)
    elif format == "json":
        # Convert tensors/arrays to JSON-serializable types
        serializable_dict = _convert_to_json_serializable(metrics)
        with open(filename, "w") as f:
            json.dump(serializable_dict, f, indent=2)
    else:
        raise ValueError(f"Unsupported save format: {format}")