Utilities

Device Management

Device management utilities for DeepCausalMMM.

This module handles: - GPU/CPU device selection - Memory management - Mixed precision training - Multi-GPU support

deepcausalmmm.utils.device.get_device(device: str | None = None) → device

Get the appropriate device for model training/inference.

Parameters:

device – Device specification (‘auto’, ‘cpu’, ‘cuda’, ‘cuda:0’, etc.) If None or ‘auto’, will use CUDA if available

Returns:

Selected device

Return type:

torch.device

deepcausalmmm.utils.device.get_amp_settings(device: device, mixed_precision: bool = True) → Tuple[GradScaler, bool]

Get Automatic Mixed Precision (AMP) settings.

Parameters:

• device – Current device

• mixed_precision – Whether to enable mixed precision training

Returns:

Tuple of (gradient scaler, use mixed precision flag)

deepcausalmmm.utils.device.move_to_device(data: Tensor | dict | list | tuple, device: device) → Tensor | dict | list | tuple

Recursively move data to specified device.

Parameters:

• data – Data to move (can be tensor, dict, list, or tuple)

• device – Target device

Returns:

Data on target device

deepcausalmmm.utils.device.clear_gpu_memory()

Clear GPU memory cache.

class deepcausalmmm.utils.device.DeviceContext(device: str | None = None, mixed_precision: bool = True)

Bases: object

Context manager for device management.

Example

with DeviceContext(device=’auto’, mixed_precision=True) as ctx:

model = model.to(ctx.device) for batch in dataloader:

with ctx.autocast():

output = model(batch)

__init__(device: str | None = None, mixed_precision: bool = True)

Initialize device context.

Parameters:

• device – Device specification

• mixed_precision – Whether to use mixed precision

class deepcausalmmm.utils.device.nullcontext(*args, **kwargs)

Bases: object

Null context manager for CPU fallback.

__init__(*args, **kwargs)

Data Generation

Config-driven synthetic data generator for DeepCausalMMM. Replaces hardcoded data generation with configurable parameters.

class deepcausalmmm.utils.data_generator.ConfigurableDataGenerator(config: Dict[str, Any] | None = None)

Bases: object

Generate synthetic MMM data using configuration parameters.

All data generation parameters are driven by configuration to ensure consistency and reproducibility across examples and tests.

__init__(config: Dict[str, Any] | None = None)

Initialize the data generator.

Parameters:

config – Configuration dictionary. If None, uses default config.

generate_mmm_dataset(n_regions: int = 2, n_weeks: int = 104, n_media_channels: int = 5, n_control_channels: int = 3) → Tuple[ndarray, ndarray, ndarray]

Generate a complete MMM dataset with realistic patterns.

Parameters:

• n_regions – Number of regions

• n_weeks – Number of weeks

• n_media_channels – Number of media channels

• n_control_channels – Number of control variables

Returns:

Tuple of (X_media, X_control, y) arrays

deepcausalmmm.utils.data_generator.get_synthetic_data_config() → Dict[str, Any]

Get default synthetic data configuration.

Returns:

Dictionary with synthetic data parameters

deepcausalmmm.utils.data_generator.update_config_with_synthetic_data(config: Dict[str, Any]) → Dict[str, Any]

Update configuration with synthetic data parameters.

Parameters:

config – Base configuration

Returns:

Updated configuration with synthetic data settings

deepcausalmmm.utils.data_generator.generate_synthetic_mmm_data(n_regions: int = 10, n_weeks: int = 52, n_media: int = 5, n_controls: int = 3, seed: int = 42)

Simple wrapper to generate synthetic MMM data as a DataFrame.

Parameters:

• n_regions – Number of regions/DMAs

• n_weeks – Number of weeks

• n_media – Number of media channels

• n_controls – Number of control variables

• seed – Random seed for reproducibility

Returns:

pandas DataFrame with synthetic MMM data