plaid.storage.common.bridge¶

def unflatten_path(key: str) -> str:
    """Unflattens a Zarr key by replacing underscores with slashes.

    Args:
        key (str): The flattened key with underscores.

    Returns:
        str: The unflattened key with slashes.
    """
    return key.replace("__", "/")

def flatten_path(key: str) -> str:
    """Flattens a path key by replacing slashes with underscores.

    Args:
        key (str): The path key to flatten.

    Returns:
        str: The flattened key with slashes replaced by underscores.
    """
    return key.replace("/", "__")

def to_sample_dict(
    var_sample_dict: dict[str, Any],
    flat_cst: dict[str, Any],
    cgns_types: dict[str, str],
    features: Optional[list[str]] = None,
) -> dict[float, dict[str, Any]]:
    """Convert variable sample dict to time-based sample dict.

    Args:
        var_sample_dict: Variable features dictionary.
        flat_cst: Constant features dictionary.
        cgns_types: CGNS types dictionary.
        features: Optional list of features to include.

    Returns:
        dict: Time-based sample dictionary.
    """
    assert not isinstance(flat_cst[next(iter(flat_cst))], dict), (
        "did you provide the complete `flat_cst` instead of the one for the considered split?"
    )

    flat_dict = var_sample_dict | flat_cst

    return flat_dict_to_sample_dict(flat_dict, cgns_types, features)

def flat_dict_to_sample_dict(
    flat_dict: dict[str, Any],
    cgns_types: dict[str, str],
    features: Optional[list[str]] = None,
) -> dict[float, dict[str, Any]]:
    """Convert a flattened sample dict into a time-indexed flat-tree mapping.

    This function processes a single flattened sample dictionary where keys
    are feature paths or feature path suffixed with '_times' (holding time
    interval specs). It splits values and their corresponding time
    specifications, slices feature arrays per interval and groups features
    by their sample time into CGNS-style flattened trees.

    Args:
        flat_dict: Mapping of flattened feature paths to values or to time
            specifications (keys ending with '_times'). Time specifications
            must be sequences or arrays reshapeable to (-1, 3) rows of
            (time, start, end).
        cgns_types: Mapping from feature path to CGNS node type; used to
            determine which None-valued features should be preserved.
        features: Optional list of feature paths to include; if None all
            features are included.

    Returns:
        A dict mapping time (float) to a flattened tree dict (path -> array or None).
    """
    if features is None:
        row_val, row_tim = _split_dict(flat_dict)
    else:  # pragma: no cover
        features_set = set(features)
        row_val, row_tim = _split_dict_feat(flat_dict, features_set)

    keys = sorted(row_tim)

    missing = set(row_val) - set(keys)
    assert not missing, f"Unexpected keys in row_val: {missing}"

    row_val = {k: row_val.get(k, None) for k in keys}
    row_tim = {k: row_tim[k] for k in keys}

    sample_flat_trees = {}
    paths_none = {}
    for (path_t, times_struc), (path_v, val) in zip(row_tim.items(), row_val.items()):
        # assert path_t == path_v, "did you forget to specify the features arg?"
        if val is None:
            assert times_struc is None
            if path_v not in paths_none and cgns_types[path_v] not in [
                "DataArray_t",
                "IndexArray_t",
            ]:
                paths_none[path_v] = None
        else:
            times_struc = np.array(times_struc, dtype=np.float64).reshape((-1, 3))
            for i, time in enumerate(times_struc[:, 0]):
                start = int(times_struc[i, 1])
                end = int(times_struc[i, 2])
                if end == -1:
                    end = None
                if val.ndim > 1:
                    values = val[:, start:end]
                else:
                    values = val[start:end]
                    if isinstance(values[0], str):
                        values = np.frombuffer(
                            values[0].encode("ascii", "strict"), dtype="|S1"
                        )
                if time in sample_flat_trees:
                    sample_flat_trees[time][path_v] = values
                else:
                    sample_flat_trees[time] = {path_v: values}

    for time, tree in sample_flat_trees.items():
        bases = {k.split("/", 1)[0] for k in tree}

        for base in bases:
            time_base = f"{base}/Time"
            tree[time_base] = np.array([1], dtype=np.int32)
            tree[f"{time_base}/IterationValues"] = np.array([1], dtype=np.int32)
            tree[f"{time_base}/TimeValues"] = np.array([time], dtype=np.float64)

        tree["CGNSLibraryVersion"] = np.array([4.0], dtype=np.float32)

        tree.update(paths_none)

    return sample_flat_trees

def to_plaid_sample(
    sample_dict: dict[float, dict[str, Any]], cgns_types: dict[str, str]
) -> Sample:
    """Convert sample dict to PLAID Sample.

    Args:
        sample_dict: Time-based sample dictionary.
        cgns_types: CGNS types dictionary.

    Returns:
        Sample: The reconstructed PLAID Sample.
    """
    sample_data = {}
    for time, flat_tree in sample_dict.items():
        sample_data[time] = unflatten_cgns_tree(flat_tree, cgns_types)

    return Sample(data=sample_data)

def plaid_to_sample_dict(
    sample: Sample,
    variable_features: Iterable[str],
    constant_features: Iterable[str],
) -> dict[str, Any]:
    """Convert PLAID Sample to sample dict.

    Args:
        sample: The PLAID Sample.
        variable_features: List of variable feature names.
        constant_features: List of constant feature names.

    Returns:
        dict[str, Any]: sample_dict
    """
    sample_dict, _, _ = build_sample_dict(sample)

    var_sample_dict = {path: sample_dict.get(path, None) for path in variable_features}
    cst_sample_dict = {path: sample_dict.get(path, None) for path in constant_features}

    return cst_sample_dict | var_sample_dict