Source code for yastn._split_combine_dict

# Copyright 2025 The YASTN Authors. All Rights Reserved.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# ==============================================================================
""" yastn.split_data_and_meta() and yastn.combine_data_and_meta() """
from __future__ import annotations

__all__ = ['split_data_and_meta', 'combine_data_and_meta']




def split_data_and_meta(d: dict, squeeze=False) -> tuple[tuple['numpy.array' | 'torch.tensor'], dict]:
    """
    Split a dictionary generated by `to_dict` methods into a tuple containing data array,
    and a dictionary `meta` with remaining information.
    Meta is a copy of `d`, where data arrays are replaced with their position in the data tuple.

    Parameters
    ----------
    d: dict
        A result of ``to_dict`` method
    squeeze: bool
        If True, and there is a single data tensor, it is unpacked from the tuple

    Return:
    data, meta
    """
    data = []  # list of datas
    meta = _split_data_and_meta(d, data)
    data = data[0] if squeeze and len(data) == 1 else tuple(data)
    return data, meta



def _split_data_and_meta(d, data):
    meta = {}
    for k in sorted(d):
        if k == "data":
            data.append(d[k])
            meta[k] = len(data) - 1
        elif isinstance(d[k], dict):
            meta[k] = _split_data_and_meta(d[k], data)
        else:
            meta[k] = d[k]
    return meta




def combine_data_and_meta(data: tuple['numpy.array' | 'torch.tensor'], meta: dict) -> dict:
    """
    Reverse :meth:`yastn.split_data_and_meta`.
    """
    d = {}
    if not isinstance(data, (list, tuple)):
        data = [data]
    for k in sorted(meta):
        if k == "data":
            d[k] = data[meta[k]]
        elif isinstance(meta[k], dict):
            d[k] = combine_data_and_meta(data, meta[k])
        else:
            d[k] = meta[k]
    return d