imt-ring 1.2.1__py3-none-any.whl

ring/algorithms/generator/base.py

@@ -0,0 +1,414 @@
+from functools import partial
+from typing import Callable, Optional, Sequence
+import warnings
+
+import jax
+import jax.numpy as jnp
+from ring import base
+from ring import utils
+from ring.algorithms import jcalc
+from ring.algorithms import kinematics
+from ring.algorithms.generator import batch
+from ring.algorithms.generator import motion_artifacts
+from ring.algorithms.generator import randomize
+from ring.algorithms.generator import transforms
+from ring.algorithms.generator import types
+import tree_utils
+
+
+class RCMG:
+    def __init__(
+        self,
+        sys: base.System | list[base.System],
+        config: jcalc.MotionConfig | list[jcalc.MotionConfig] = jcalc.MotionConfig(),
+        setup_fn: Optional[types.SETUP_FN] = None,
+        finalize_fn: Optional[types.FINALIZE_FN] = None,
+        add_X_imus: bool = False,
+        add_X_imus_kwargs: Optional[dict] = None,
+        add_X_jointaxes: bool = False,
+        add_X_jointaxes_kwargs: Optional[dict] = None,
+        add_y_relpose: bool = False,
+        add_y_rootincl: bool = False,
+        sys_ml: Optional[base.System] = None,
+        randomize_positions: bool = False,
+        randomize_motion_artifacts: bool = False,
+        randomize_joint_params: bool = False,
+        randomize_anchors: bool = False,
+        randomize_anchors_kwargs: Optional[dict] = None,
+        randomize_hz: bool = False,
+        randomize_hz_kwargs: Optional[dict] = None,
+        imu_motion_artifacts: bool = False,
+        imu_motion_artifacts_kwargs: Optional[dict] = None,
+        dynamic_simulation: bool = False,
+        dynamic_simulation_kwargs: Optional[dict] = None,
+        output_transform: Optional[Callable] = None,
+        keep_output_extras: bool = False,
+        use_link_number_in_Xy: bool = False,
+    ) -> None:
+
+        randomize_anchors_kwargs = _copy_kwargs(randomize_anchors_kwargs)
+        randomize_hz_kwargs = _copy_kwargs(randomize_hz_kwargs)
+
+        if randomize_hz:
+            finalize_fn = randomize.randomize_hz_finalize_fn_factory(finalize_fn)
+
+        partial_build_gen = partial(
+            _build_generator_lazy,
+            setup_fn=setup_fn,
+            finalize_fn=finalize_fn,
+            add_X_imus=add_X_imus,
+            add_X_imus_kwargs=add_X_imus_kwargs,
+            add_X_jointaxes=add_X_jointaxes,
+            add_X_jointaxes_kwargs=add_X_jointaxes_kwargs,
+            add_y_relpose=add_y_relpose,
+            add_y_rootincl=add_y_rootincl,
+            randomize_positions=randomize_positions,
+            randomize_motion_artifacts=randomize_motion_artifacts,
+            randomize_joint_params=randomize_joint_params,
+            imu_motion_artifacts=imu_motion_artifacts,
+            imu_motion_artifacts_kwargs=imu_motion_artifacts_kwargs,
+            dynamic_simulation=dynamic_simulation,
+            dynamic_simulation_kwargs=dynamic_simulation_kwargs,
+            output_transform=output_transform,
+            keep_output_extras=keep_output_extras,
+            use_link_number_in_Xy=use_link_number_in_Xy,
+        )
+
+        sys, config = utils.to_list(sys), utils.to_list(config)
+
+        if randomize_anchors:
+            assert (
+                len(sys) == 1
+            ), "If `randomize_anchors`, then only one system is expected"
+            sys = randomize.randomize_anchors(sys[0], **randomize_anchors_kwargs)
+
+        zip_sys_config = False
+        if randomize_hz:
+            zip_sys_config = True
+            sys, config = randomize.randomize_hz(sys, config, **randomize_hz_kwargs)
+
+        if sys_ml is None:
+            # TODO
+            if False and len(sys) > 1:
+                warnings.warn(
+                    "Batched simulation with multiple systems but no explicit `sys_ml`"
+                )
+            sys_ml = sys[0]
+
+        self.gens = []
+        if zip_sys_config:
+            for _sys, _config in zip(sys, config):
+                self.gens.append(
+                    partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml)
+                )
+        else:
+            for _sys in sys:
+                for _config in config:
+                    self.gens.append(
+                        partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml)
+                    )
+
+    def _to_data(self, sizes, seed, jit):
+        return batch.batch_generators_eager_to_list(
+            self.gens, sizes, seed=seed, jit=jit
+        )
+
+    def to_list(self, sizes: int | list[int] = 1, seed: int = 1, jit: bool = False):
+        return self._to_data(sizes, seed, jit)
+
+    def to_pickle(
+        self,
+        path: str,
+        sizes: int | list[int] = 1,
+        seed: int = 1,
+        jit: bool = False,
+        overwrite: bool = True,
+    ) -> None:
+        data = tree_utils.tree_batch(self._to_data(sizes, seed, jit))
+        utils.pickle_save(data, path, overwrite=overwrite)
+
+    def to_hdf5(
+        self,
+        path: str,
+        sizes: int | list[int] = 1,
+        seed: int = 1,
+        jit: bool = False,
+        overwrite: bool = True,
+    ) -> None:
+        data = tree_utils.tree_batch(self._to_data(sizes, seed, jit))
+        utils.hdf5_save(path, data, overwrite=overwrite)
+
+    def to_eager_gen(
+        self,
+        batchsize: int = 1,
+        sizes: int | list[int] = 1,
+        seed: int = 1,
+        jit: bool = False,
+    ) -> types.BatchedGenerator:
+        return batch.batch_generators_eager(
+            self.gens, sizes, batchsize, seed=seed, jit=jit
+        )
+
+    def to_lazy_gen(
+        self, sizes: int | list[int] = 1, jit: bool = True
+    ) -> types.BatchedGenerator:
+        return batch.batch_generators_lazy(self.gens, sizes, jit=jit)
+
+    @staticmethod
+    def eager_gen_from_paths(
+        paths: str | list[str],
+        batchsize: int,
+        include_samples: Optional[list[int]] = None,
+        shuffle: bool = True,
+        load_all_into_memory: bool = False,
+        tree_transform=None,
+    ) -> tuple[types.BatchedGenerator, int]:
+        paths = utils.to_list(paths)
+        return batch.batched_generator_from_paths(
+            paths,
+            batchsize,
+            include_samples,
+            shuffle,
+            load_all_into_memory=load_all_into_memory,
+            tree_transform=tree_transform,
+        )
+
+
+def _copy_kwargs(kwargs: dict | None) -> dict:
+    return dict() if kwargs is None else kwargs.copy()
+
+
+def _build_generator_lazy(
+    sys: base.System,
+    config: jcalc.MotionConfig,
+    setup_fn: types.SETUP_FN | None,
+    finalize_fn: types.FINALIZE_FN | None,
+    add_X_imus: bool,
+    add_X_imus_kwargs: dict | None,
+    add_X_jointaxes: bool,
+    add_X_jointaxes_kwargs: dict | None,
+    add_y_relpose: bool,
+    add_y_rootincl: bool,
+    sys_ml: base.System,
+    randomize_positions: bool,
+    randomize_motion_artifacts: bool,
+    randomize_joint_params: bool,
+    imu_motion_artifacts: bool,
+    imu_motion_artifacts_kwargs: dict | None,
+    dynamic_simulation: bool,
+    dynamic_simulation_kwargs: dict | None,
+    output_transform: Callable | None,
+    keep_output_extras: bool,
+    use_link_number_in_Xy: bool,
+) -> types.Generator | types.GeneratorWithOutputExtras:
+    assert config.is_feasible()
+
+    imu_motion_artifacts_kwargs = _copy_kwargs(imu_motion_artifacts_kwargs)
+    dynamic_simulation_kwargs = _copy_kwargs(dynamic_simulation_kwargs)
+    add_X_imus_kwargs = _copy_kwargs(add_X_imus_kwargs)
+    add_X_jointaxes_kwargs = _copy_kwargs(add_X_jointaxes_kwargs)
+
+    # default kwargs values
+    if "hide_injected_bodies" not in imu_motion_artifacts_kwargs:
+        imu_motion_artifacts_kwargs["hide_injected_bodies"] = True
+
+    if add_X_jointaxes or add_y_relpose or add_y_rootincl:
+        if len(sys_ml.findall_imus()) > 0:
+            # warnings.warn("Automatically removed the IMUs from `sys_ml`.")
+            sys_noimu, _ = sys_ml.make_sys_noimu()
+        else:
+            sys_noimu = sys_ml
+
+    unactuated_subsystems = []
+    if imu_motion_artifacts:
+        assert dynamic_simulation
+        unactuated_subsystems = motion_artifacts.unactuated_subsystem(sys)
+        sys = motion_artifacts.inject_subsystems(sys, **imu_motion_artifacts_kwargs)
+        assert "unactuated_subsystems" not in dynamic_simulation_kwargs
+        dynamic_simulation_kwargs["unactuated_subsystems"] = unactuated_subsystems
+
+        if not randomize_motion_artifacts:
+            warnings.warn(
+                "`imu_motion_artifacts` is enabled but not `randomize_motion_artifacts`"
+            )
+
+        if "hide_injected_bodies" in imu_motion_artifacts_kwargs:
+            if imu_motion_artifacts_kwargs["hide_injected_bodies"] and False:
+                warnings.warn(
+                    "The flag `hide_injected_bodies` in `imu_motion_artifacts_kwargs` "
+                    "is set. This will try to hide injected bodies. This feature is "
+                    "experimental."
+                )
+
+        if "prob_rigid" in imu_motion_artifacts_kwargs:
+            assert randomize_motion_artifacts, (
+                "`prob_rigid` works by overwriting damping and stiffness parameters "
+                "using the `randomize_motion_artifacts` flag, so it must be enabled."
+            )
+
+    noop = lambda gen: gen
+    return GeneratorPipe(
+        transforms.GeneratorTrafoSetupFn(setup_fn) if setup_fn is not None else noop,
+        (
+            transforms.GeneratorTrafoSetupFn(jcalc._init_joint_params)
+            if randomize_joint_params
+            else noop
+        ),
+        transforms.GeneratorTrafoRandomizePositions() if randomize_positions else noop,
+        (
+            transforms.GeneratorTrafoSetupFn(
+                motion_artifacts.setup_fn_randomize_damping_stiffness_factory(
+                    prob_rigid=imu_motion_artifacts_kwargs.get("prob_rigid", 0.0),
+                    all_imus_either_rigid_or_flex=imu_motion_artifacts_kwargs.get(
+                        "all_imus_either_rigid_or_flex", False
+                    ),
+                    imus_surely_rigid=imu_motion_artifacts_kwargs.get(
+                        "imus_surely_rigid", []
+                    ),
+                )
+            )
+            if (imu_motion_artifacts and randomize_motion_artifacts)
+            else noop
+        ),
+        # all the generator trafors before this point execute in reverse order
+        # to see this, consider gen[0] and gen[1]
+        # the GeneratorPipe will unpack into the following:
+        # gen[1] will unfold into
+        # >>> sys = gen[1].setup_fn(sys)
+        # >>> return gen[0](sys)
+        # <-------------------- GENERATOR MIDDLE POINT ------------------------->
+        # all the generator trafos after this point execute in order
+        # >>> Xy, extras = gen[-2](*args)
+        # >>> return gen[-1].finalize_fn(extras)
+        (
+            transforms.GeneratorTrafoDynamicalSimulation(**dynamic_simulation_kwargs)
+            if dynamic_simulation
+            else noop
+        ),
+        (
+            motion_artifacts.GeneratorTrafoHideInjectedBodies()
+            if (
+                imu_motion_artifacts
+                and imu_motion_artifacts_kwargs["hide_injected_bodies"]
+            )
+            else noop
+        ),
+        (
+            transforms.GeneratorTrafoFinalizeFn(finalize_fn)
+            if finalize_fn is not None
+            else noop
+        ),
+        transforms.GeneratorTrafoIMU(**add_X_imus_kwargs) if add_X_imus else noop,
+        (
+            transforms.GeneratorTrafoJointAxisSensor(
+                sys_noimu, **add_X_jointaxes_kwargs
+            )
+            if add_X_jointaxes
+            else noop
+        ),
+        transforms.GeneratorTrafoRelPose(sys_noimu) if add_y_relpose else noop,
+        transforms.GeneratorTrafoRootIncl(sys_noimu) if add_y_rootincl else noop,
+        (
+            transforms.GeneratorTrafoNames2Indices(sys_noimu)
+            if use_link_number_in_Xy
+            else noop
+        ),
+        GeneratorTrafoRemoveInputExtras(sys),
+        noop if keep_output_extras else GeneratorTrafoRemoveOutputExtras(),
+        (
+            transforms.GeneratorTrafoLambda(output_transform, input=False)
+            if output_transform is not None
+            else noop
+        ),
+    )(config)
+
+
+def _generator_with_extras(
+    config: jcalc.MotionConfig,
+) -> types.GeneratorWithInputOutputExtras:
+    def generator(
+        key: types.PRNGKey, sys: base.System
+    ) -> tuple[types.Xy, types.OutputExtras]:
+        if config.cor:
+            sys = sys._replace_free_with_cor()
+
+        key_start = key
+        # build generalized coordintes vector `q`
+        q_list = []
+
+        def draw_q(key, __, link_type, link):
+            joint_params = link.joint_params
+            # limit scope
+            joint_params = (
+                joint_params[link_type]
+                if link_type in joint_params
+                else joint_params["default"]
+            )
+            if key is None:
+                key = key_start
+            key, key_t, key_value = jax.random.split(key, 3)
+            draw_fn = jcalc.get_joint_model(link_type).rcmg_draw_fn
+            if draw_fn is None:
+                raise Exception(f"The joint type {link_type} has no draw fn specified.")
+            q_link = draw_fn(config, key_t, key_value, sys.dt, joint_params)
+            # even revolute and prismatic joints must be 2d arrays
+            q_link = q_link if q_link.ndim == 2 else q_link[:, None]
+            q_list.append(q_link)
+            return key
+
+        keys = sys.scan(draw_q, "ll", sys.link_types, sys.links)
+        # stack of keys; only the last key is unused
+        key = keys[-1]
+
+        q = jnp.concatenate(q_list, axis=1)
+
+        # do forward kinematics
+        x, _ = jax.vmap(kinematics.forward_kinematics_transforms, (None, 0))(sys, q)
+
+        Xy = ({}, {})
+        return Xy, (key, q, x, sys)
+
+    return generator
+
+
+class GeneratorPipe:
+    def __init__(self, *gen_trafos: Sequence[types.GeneratorTrafo]):
+        self._gen_trafos = gen_trafos
+
+    def __call__(
+        self, config: jcalc.MotionConfig
+    ) -> (
+        types.GeneratorWithInputOutputExtras
+        | types.GeneratorWithOutputExtras
+        | types.GeneratorWithInputExtras
+        | types.Generator
+    ):
+        gen = _generator_with_extras(config)
+        for trafo in self._gen_trafos:
+            gen = trafo(gen)
+        return gen
+
+
+class GeneratorTrafoRemoveInputExtras(types.GeneratorTrafo):
+    def __init__(self, sys: base.System):
+        self.sys = sys
+
+    def __call__(
+        self,
+        gen: types.GeneratorWithInputExtras | types.GeneratorWithInputOutputExtras,
+    ) -> types.Generator | types.GeneratorWithOutputExtras:
+        def _gen(key):
+            return gen(key, self.sys)
+
+        return _gen
+
+
+class GeneratorTrafoRemoveOutputExtras(types.GeneratorTrafo):
+    def __call__(
+        self,
+        gen: types.GeneratorWithOutputExtras | types.GeneratorWithInputOutputExtras,
+    ) -> types.Generator | types.GeneratorWithInputExtras:
+        def _gen(*args):
+            return gen(*args)[0]
+
+        return _gen

ring/algorithms/generator/batch.py

@@ -0,0 +1,282 @@
+from pathlib import Path
+import random
+from typing import Optional
+import warnings
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from ring import utils
+from ring.algorithms.generator import types
+from tqdm import tqdm
+import tree_utils
+from tree_utils import tree_batch
+
+
+def _build_batch_matrix(batchsizes: list[int]) -> jax.Array:
+    arr = []
+    for i, l in enumerate(batchsizes):
+        arr += [i] * l
+    return jnp.array(arr)
+
+
+def batch_generators_lazy(
+    generators: types.Generator | list[types.Generator],
+    batchsizes: int | list[int] = 1,
+    jit: bool = True,
+) -> types.BatchedGenerator:
+    """Create a large generator by stacking multiple generators lazily."""
+    generators = utils.to_list(generators)
+
+    generators, batchsizes = _process_sizes_batchsizes_generators(
+        generators, batchsizes
+    )
+
+    batch_arr = _build_batch_matrix(batchsizes)
+    bs_total = len(batch_arr)
+    pmap, vmap = utils.distribute_batchsize(bs_total)
+    batch_arr = batch_arr.reshape((pmap, vmap))
+
+    pmap_trafo = jax.pmap
+    # single GPU node, then do jit + vmap instead of pmap
+    # this allows e.g. better NAN debugging capabilities
+    if pmap == 1:
+        pmap_trafo = lambda f: jax.jit(jax.vmap(f))
+    if not jit:
+        pmap_trafo = lambda f: jax.vmap(f)
+
+    @pmap_trafo
+    @jax.vmap
+    def _generator(key, which_gen: int):
+        return jax.lax.switch(which_gen, generators, key)
+
+    def generator(key):
+        pmap_vmap_keys = jax.random.split(key, bs_total).reshape((pmap, vmap, 2))
+        data = _generator(pmap_vmap_keys, batch_arr)
+
+        # merge pmap and vmap axis
+        data = utils.merge_batchsize(data, pmap, vmap)
+        return data
+
+    return generator
+
+
+def batch_generators_eager_to_list(
+    generators: types.Generator | list[types.Generator],
+    sizes: int | list[int],
+    seed: int = 1,
+    jit: bool = True,
+) -> list[tree_utils.PyTree]:
+    "Returns list of unbatched sequences as numpy arrays."
+    generators, sizes = _process_sizes_batchsizes_generators(generators, sizes)
+
+    key = jax.random.PRNGKey(seed)
+    data = []
+    for gen, size in tqdm(zip(generators, sizes), desc="eager data generation"):
+        key, consume = jax.random.split(key)
+        sample = batch_generators_lazy(gen, size, jit=jit)(consume)
+        # converts also to numpy
+        sample = jax.device_get(sample)
+        data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
+    return data
+
+
+def _is_nan(ele: tree_utils.PyTree, i: int, verbose: bool = False):
+    isnan = np.any([np.any(np.isnan(arr)) for arr in jax.tree_util.tree_leaves(ele)])
+    if isnan:
+        X, y = ele
+        dt = X["dt"].flatten()[0]
+        if verbose:
+            print(f"Sample with idx={i} is nan. It will be replaced. (dt={dt})")
+        return True
+    return False
+
+
+def _replace_elements_w_nans(list_of_data: list, include_samples: list[int]) -> list:
+    list_of_data_nonan = []
+    for i, ele in enumerate(list_of_data):
+        if _is_nan(ele, i, verbose=True):
+            while True:
+                j = random.choice(include_samples)
+                if not _is_nan(list_of_data[j], j):
+                    ele = list_of_data[j]
+                    break
+        list_of_data_nonan.append(ele)
+    return list_of_data_nonan
+
+
+_list_of_data = None
+_paths = None
+
+
+def _data_fn_from_paths(
+    paths: list[str],
+    include_samples: list[int] | None,
+    load_all_into_memory: bool,
+    tree_transform,
+):
+    "`data_fn` returns numpy arrays."
+    global _list_of_data, _paths
+
+    # expanduser
+    paths = [utils.parse_path(p, mkdir=False) for p in paths]
+
+    extensions = list(set([Path(p).suffix for p in paths]))
+    assert len(extensions) == 1
+
+    if extensions[0] == ".h5":
+        N = sum([utils.hdf5_load_length(p) for p in paths])
+
+    if extensions[0] == ".h5" and not load_all_into_memory:
+
+        def data_fn(indices: list[int]):
+            tree = utils.hdf5_load_from_multiple(paths, indices)
+            return tree if tree_transform is None else tree_transform(tree)
+
+    else:
+
+        if extensions[0] == ".h5":
+            load_from_path = utils.hdf5_load
+        else:
+            load_from_path = utils.pickle_load
+
+        def load_fn(path):
+            tree = load_from_path(path)
+            tree = tree if tree_transform is None else tree_transform(tree)
+            return [
+                jax.tree_map(lambda arr: arr[i], tree)
+                for i in range(tree_utils.tree_shape(tree))
+            ]
+
+        if paths != _paths or len(_list_of_data) == 0:
+            _paths = paths
+
+            _list_of_data = []
+            for p in paths:
+                _list_of_data += load_fn(p)
+
+        N = len(_list_of_data)
+
+        list_of_data = _replace_elements_w_nans(_list_of_data, include_samples)
+
+        if include_samples is not None:
+            list_of_data = [
+                ele if i in include_samples else None
+                for i, ele in enumerate(list_of_data)
+            ]
+
+        def data_fn(indices: list[int]):
+            return tree_batch([list_of_data[i] for i in indices], backend="numpy")
+
+    if include_samples is None:
+        include_samples = list(range(N))
+
+    return data_fn, include_samples.copy()
+
+
+def _generator_from_data_fn(
+    data_fn,
+    include_samples: list[int],
+    shuffle: bool,
+    batchsize: int,
+) -> types.BatchedGenerator:
+    # such that we don't mutate out of scope
+    include_samples = include_samples.copy()
+
+    N = len(include_samples)
+    n_batches, i = N // batchsize, 0
+
+    def generator(key: jax.Array):
+        nonlocal i
+        if shuffle and i == 0:
+            random.shuffle(include_samples)
+
+        start, stop = i * batchsize, (i + 1) * batchsize
+        batch = data_fn(include_samples[start:stop])
+
+        i = (i + 1) % n_batches
+        return utils.pytree_deepcopy(batch)
+
+    return generator
+
+
+def batched_generator_from_paths(
+    paths: list[str],
+    batchsize: int,
+    include_samples: Optional[list[int]] = None,
+    shuffle: bool = True,
+    load_all_into_memory: bool = False,
+    tree_transform=None,
+) -> tuple[types.BatchedGenerator, int]:
+    "Returns: gen, where gen(key) -> Pytree[numpy]"
+    data_fn, include_samples = _data_fn_from_paths(
+        paths, include_samples, load_all_into_memory, tree_transform
+    )
+
+    N = len(include_samples)
+    assert N >= batchsize
+
+    generator = _generator_from_data_fn(data_fn, include_samples, shuffle, batchsize)
+
+    return generator, N
+
+
+def batched_generator_from_list(
+    data: list,
+    batchsize: int,
+    shuffle: bool = True,
+    drop_last: bool = True,
+) -> types.BatchedGenerator:
+    assert drop_last, "Not `drop_last` is currently not implemented."
+    assert len(data) >= batchsize
+
+    def data_fn(indices: list[int]):
+        return tree_batch([data[i] for i in indices])
+
+    return _generator_from_data_fn(data_fn, list(range(len(data))), shuffle, batchsize)
+
+
+def batch_generators_eager(
+    generators: types.Generator | list[types.Generator],
+    sizes: int | list[int],
+    batchsize: int,
+    shuffle: bool = True,
+    drop_last: bool = True,
+    seed: int = 1,
+    jit: bool = True,
+) -> types.BatchedGenerator:
+    """Eagerly create a large precomputed generator by calling multiple generators
+    and stacking their output."""
+
+    data = batch_generators_eager_to_list(generators, sizes, seed=seed, jit=jit)
+    return batched_generator_from_list(data, batchsize, shuffle, drop_last)
+
+
+def _process_sizes_batchsizes_generators(
+    generators: types.Generator | list[types.Generator],
+    batchsizes_or_sizes: int | list[int],
+) -> tuple[list, list]:
+    generators = utils.to_list(generators)
+    assert len(generators) > 0, "No generator was passed."
+
+    if isinstance(batchsizes_or_sizes, int):
+        assert (
+            batchsizes_or_sizes // len(generators)
+        ) > 0, f"Batchsize or size too small. {batchsizes_or_sizes} < {len(generators)}"
+        list_sizes = len(generators) * [batchsizes_or_sizes // len(generators)]
+    else:
+        list_sizes = batchsizes_or_sizes
+        assert 0 not in list_sizes
+
+    assert len(generators) == len(list_sizes)
+
+    _WARN_SIZE = 4096
+    for size in list_sizes:
+        if size >= _WARN_SIZE:
+            warnings.warn(
+                f"A generator will be called with a large batchsize of {size} "
+                f"(warn limit is {_WARN_SIZE}). The generator sizes are {list_sizes}."
+            )
+            break
+
+    return generators, list_sizes