from ..base import BatchMatrixData
from .diffusion import DiffusionGenerator


class DiffusionAdvectionGenerator:

    def __init__(self, diffusion_generator: DiffusionGenerator):
        self.diffusion_generator = diffusion_generator

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        return self.diffusion_generator.generate_advection_batch(batch_size, num_context_pairs)
