"""
Factory tools for managing random draw generation in Biogeme.
This module defines classes for encapsulating and orchestrating
the creation of random draws used in simulation-based estimation.
It distinguishes between native and user-defined draw generators,
validates inputs, and constructs a final tensor of draws with shape:
(sample_size, number_of_draws, number_of_variables)
Michel Bierlaire
Wed Mar 26 19:30:36 2025
"""
from dataclasses import dataclass
from collections.abc import Callable
import numpy as np
from biogeme.exceptions import BiogemeError
from .native_draws import (
native_random_number_generators,
RandomNumberGeneratorTuple,
)
[docs]
@dataclass
class DrawSpec:
"""
Encapsulates the configuration for generating draws for a specific variable.
:param name: Name of the variable that requires simulated draws.
:param draw_type: Identifier for the type of draw (native or user-defined).
:param generator: A callable that takes (sample_size, number_of_draws)
and returns a NumPy array of draws.
"""
name: str
draw_type: str
generator: Callable[[int, int], np.ndarray]
[docs]
class DrawFactory:
"""
Manages native and user-defined random draw generators and builds
draw specifications and arrays for multiple variables.
This class is useful for transforming a mapping of variable names and draw types
into a single array of random numbers ready for use in simulation-based estimation.
"""
def __init__(self, user_generators: dict[str, RandomNumberGeneratorTuple] | None):
"""
Initializes the manager and validates user-defined generators.
:param user_generators: A dictionary of user-defined draw types.
:raises ValueError: if any user generator name collides with native keywords.
"""
self.native_generators = native_random_number_generators
self.user_generators = {} if user_generators is None else user_generators
self._validate_reserved_keywords()
def _validate_reserved_keywords(self) -> None:
"""Ensure no user-defined generator overrides a native one."""
for key in self.native_generators:
if key in self.user_generators:
raise ValueError(
f"{key} is a reserved keyword and cannot be used "
"as a user-defined draw generator name."
)
[docs]
def get_generator(
self, draw_type: str, name: str
) -> Callable[[int, int], np.ndarray]:
"""
Retrieves a draw generator function based on the requested type.
:param draw_type: Type identifier of the draw (e.g., 'UNIFORM', 'HALTON').
:param name: Variable name (used for error context).
:return: A function that generates a NumPy array of draws.
:raises BiogemeError: if the draw type is not recognized.
"""
if draw_type in self.native_generators:
return self.native_generators[draw_type].generator
if draw_type in self.user_generators:
return self.user_generators[draw_type].generator
raise BiogemeError(
f"Unknown draw type '{draw_type}' for variable '{name}'. "
f"Available native types: {list(self.native_generators.keys())}. "
f"User-defined types: {list(self.user_generators.keys())}."
)
[docs]
def make_draw_specs(
self, draw_types: dict[str, str], variable_names: list[str]
) -> list[DrawSpec]:
"""
Generates a list of DrawSpec objects for each variable.
:param draw_types: Mapping from variable name to draw type.
:param variable_names: List of variable names requiring simulated draws.
:return: List of fully constructed DrawSpec objects.
"""
return [
DrawSpec(
name=name,
draw_type=draw_types[name],
generator=self.get_generator(draw_types[name], name),
)
for name in variable_names
]
[docs]
def generate_draws(
self,
draw_types: dict[str, str],
variable_names: list[str],
sample_size: int,
number_of_draws: int,
) -> np.ndarray:
"""
Generates a 3D NumPy array of draws for all specified variables.
:param draw_types: Mapping from variable name to draw type.
:param variable_names: Ordered list of variable names.
:param sample_size: Number of observations in the sample.
:param number_of_draws: Number of Monte Carlo draws per observation.
:return: A NumPy array of shape (sample_size, number_of_draws, len(variable_names)).
:raises BiogemeError: if any generator returns a mis-shaped array.
"""
specs = self.make_draw_specs(draw_types, variable_names)
draws = []
for spec in specs:
array = spec.generator(sample_size, number_of_draws)
self._check_shape(spec.name, array, sample_size, number_of_draws)
draws.append(array)
return np.moveaxis(np.array(draws), 0, -1)
def _check_shape(self, name, array, sample_size, number_of_draws):
"""
Verifies that the shape of a generated draw array matches expectations.
:param name: Variable name.
:param array: Array returned by the draw generator.
:param sample_size: Expected number of rows.
:param number_of_draws: Expected number of columns.
:raises BiogemeError: if the shape is invalid.
"""
if array.shape != (sample_size, number_of_draws):
raise BiogemeError(
f"Draws for '{name}' must have shape ({sample_size}, {number_of_draws}), "
f"but got {array.shape}."
)