Compression Base Classes

@abstractmethod
def compress(
    self,
    values: dict[str, jax.Array],
    state: RecordingState,
    key: jax.Array,
) -> tuple[
    dict[str, jax.Array],  # compressed data
    RecordingState,  # updated recording state
]:
    """Compress field values at the current time step.

    Args:
        values (dict[str, jax.Array]): Dictionary mapping field names to their values.
        state (RecordingState): Current recording state.
        key (jax.Array): Random key for stochastic operations.

    Returns:
        tuple[dict[str, jax.Array], RecordingState]: Tuple containing:
            - Dictionary of compressed field values
            - Updated recording state
    """
    del values, state, key
    raise NotImplementedError()

@abstractmethod
def decompress(
    self,
    values: dict[str, jax.Array],
    state: RecordingState,
    key: jax.Array,
) -> dict[str, jax.Array]:
    """Decompress field values back to their original form.

    Args:
        values (dict[str, jax.Array]): Dictionary mapping field names to their compressed values.
        state (RecordingState): Current recording state.
        key (jax.Array): Random key for stochastic operations.

    Returns:
        dict[str, jax.Array]: Dictionary mapping field names to their decompressed values.
    """
    del (
        values,
        state,
        key,
    )
    raise NotImplementedError()

@abstractmethod
def init_shapes(
    self,
    input_shape_dtypes: dict[str, jax.ShapeDtypeStruct],
) -> tuple[
    Self,
    dict[str, jax.ShapeDtypeStruct],  # data
    dict[str, jax.ShapeDtypeStruct],  # state shapes/dtypes
]:
    """Initialize shapes and sizes for the compression module.

    Args:
        input_shape_dtypes (dict[str, jax.ShapeDtypeStruct]): Dictionary mapping field names to their input
            shapes/dtypes.

    Returns:
        tuple[Self, dict[str, jax.ShapeDtypeStruct], dict[str, jax.ShapeDtypeStruct]]: Tuple containing:
            - Self: Updated instance of the compression module
            - Dictionary mapping field names to their output shapes/dtypes
            - Dictionary mapping field names to their state shapes/dtypes
    """
    del input_shape_dtypes
    raise NotImplementedError()

@abstractmethod
def compress(
    self,
    values: dict[str, jax.Array],
    state: RecordingState,
    time_idx: jax.Array,  # scalar
    key: jax.Array,
) -> tuple[
    dict[str, jax.Array],
    RecordingState,  # updated recording state
]:
    """Compress field values at a given time step.

    Args:
        values (dict[str, jax.Array]): Dictionary mapping field names to their values.
        state (RecordingState): Current recording state.
        time_idx (jax.Array): Current time step index.
        key (jax.Array): Random key for stochastic operations.

    Returns:
        tuple[dict[str, jax.Array], RecordingState]: Tuple containing:
            - Dictionary of compressed field values
            - Updated recording state
    """
    del values, state, time_idx, key
    raise NotImplementedError()

@abstractmethod
def decompress(
    self,
    values: list[dict[str, jax.Array]],  # array values requested above
    state: RecordingState,
    arr_indices: jax.Array,
    time_idx: jax.Array,  # scalar
    key: jax.Array,
) -> dict[str, jax.Array]:
    """Decompress field values to reconstruct data for a time step.

    Args:
        values (list[dict[str, jax.Array]]): List of dictionaries containing array values needed for reconstruction.
        state (RecordingState): Current recording state.
        arr_indices (jax.Array): Array indices needed for reconstruction.
        time_idx (jax.Array): Time step index to reconstruct. scalar value.
        key (jax.Array): Random key for stochastic operations.

    Returns:
        dict[str, jax.Array]: Dictionary of reconstructed field values.
    """
    del values, state, arr_indices, time_idx, key
    raise NotImplementedError()

@abstractmethod
def indices_to_decompress(
    self,
    time_idx: jax.Array,  # scalar
) -> jax.Array:  # 1d-list of array indices necessary to reconstruct
    """Get array indices needed to reconstruct data for a given time step.

    Args:
        time_idx (jax.Array): Time step index to reconstruct.

    Returns:
        jax.Array: Array of indices needed to reconstruct the data for this time step.
    """
    del time_idx
    raise NotImplementedError()

@abstractmethod
def init_shapes(
    self,
    input_shape_dtypes: dict[str, jax.ShapeDtypeStruct],
    time_steps_max: int,  # maximum number of time steps
) -> tuple[
    Self,
    int,  # array size (number of latent time steps)
    dict[str, jax.ShapeDtypeStruct],  # data
    dict[str, jax.ShapeDtypeStruct],  # state shapes
]:
    """Initialize shapes and sizes for the time step filter.

    Args:
        input_shape_dtypes (dict[str, jax.ShapeDtypeStruct]): Dictionary mapping field names to their
            shape/dtype information.
        time_steps_max (int): Maximum number of time steps in the simulation.

    Returns:
        tuple[Self, int, dict[str, jax.ShapeDtypeStruct], dict[str, jax.ShapeDtypeStruct]]: A tuple containing:
            - Updated filter instance
            - Size of array for storing filtered data
            - Dictionary of data shapes/dtypes
            - Dictionary of state shapes/dtypes
    """
    del input_shape_dtypes, time_steps_max
    raise NotImplementedError()

@abstractmethod
def time_to_array_index(
    self,
    time_idx: int,  # scalar
) -> int:  # array index if not filtered, else -1
    """Convert a time step index to its corresponding array index.

    Args:
        time_idx (int): Time step index to convert.

    Returns:
        int: The corresponding array index if the time step is not filtered,
            or -1 if the time step is filtered out.
    """
    del time_idx
    raise NotImplementedError()