data-designer 0.3.8rc1__py3-none-any.whl → 0.4.0__py3-none-any.whl

data_designer/engine/sampling_gen/constraints.py

@@ -1,100 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
-
-from numpy.typing import NDArray
-
-from data_designer.config.base import ConfigBase
-from data_designer.config.sampler_constraints import (
-    ColumnInequalityConstraint,
-    Constraint,
-    ConstraintType,
-    InequalityOperator,
-    ScalarInequalityConstraint,
-)
-from data_designer.lazy_heavy_imports import np, pd
-
-if TYPE_CHECKING:
-    import numpy as np
-    import pandas as pd
-
-
-class ConstraintChecker(ConfigBase, ABC):
-    constraint: Constraint
-
-    def get_required_column_names(self) -> tuple[str, ...]:
-        return (self.constraint.target_column,)
-
-    @abstractmethod
-    def check(self, dataframe: pd.DataFrame) -> NDArray[np.bool_]: ...
-
-
-class WithCompareMixin:
-    @property
-    def lhs(self) -> str:
-        return self.constraint.target_column
-
-    def compare(self, lhs: float | int | NDArray, rhs: float | int | NDArray) -> bool | NDArray[np.bool_]:
-        operator = {
-            InequalityOperator.LT: np.less,
-            InequalityOperator.LE: np.less_equal,
-            InequalityOperator.GT: np.greater,
-            InequalityOperator.GE: np.greater_equal,
-        }[InequalityOperator(self.constraint.operator)]
-        return operator(lhs, rhs)
-
-
-class ScalarInequalityChecker(ConstraintChecker, WithCompareMixin):
-    """Compare a column to a scalar value.
-
-    Args:
-        column_name: Name of the constrained column. Will be
-            used as the left-hand side (lhs) of the comparison.
-        operator: Comparison operator.
-        rhs: Scalar value to compare against.
-    """
-
-    constraint: ScalarInequalityConstraint
-
-    def check(self, dataframe: pd.DataFrame) -> NDArray[np.bool_]:
-        return self.compare(dataframe[self.lhs].values, self.constraint.rhs)
-
-
-class ColumnInequalityChecker(ConstraintChecker, WithCompareMixin):
-    """Compare the values of two columns.
-
-    Args:
-        column_name: Name of the constrained column. Will be
-            used as the left-hand side (lhs) of the comparison.
-        operator: Comparison operator.
-        rhs: Name of the column to compare against.
-    """
-
-    constraint: ColumnInequalityConstraint
-
-    def get_required_column_names(self) -> tuple[str, ...]:
-        """Return the names of columns required for the constraint.
-
-        Note that order matters. Edges in the DAG are created as column_names[1], column_names[0].
-        """
-        return (self.lhs, self.constraint.rhs)
-
-    def check(self, dataframe: pd.DataFrame) -> NDArray[np.bool_]:
-        return self.compare(
-            dataframe[self.lhs].values,
-            dataframe[self.constraint.rhs].values,
-        )
-
-
-CONSTRAINT_TYPE_TO_CHECKER = {
-    ConstraintType.SCALAR_INEQUALITY: ScalarInequalityChecker,
-    ConstraintType.COLUMN_INEQUALITY: ColumnInequalityChecker,
-}
-
-
-def get_constraint_checker(constraint_type: ConstraintType) -> type[ConstraintChecker]:
-    return CONSTRAINT_TYPE_TO_CHECKER[ConstraintType(constraint_type)]

data_designer/engine/sampling_gen/data_sources/base.py

@@ -1,217 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
-
-from numpy.typing import NDArray
-
-from data_designer.config.sampler_params import SamplerParamsT
-from data_designer.engine.sampling_gen.utils import check_random_state
-from data_designer.lazy_heavy_imports import np, pd, scipy
-
-if TYPE_CHECKING:
-    import numpy as np
-    import pandas as pd
-    import scipy
-
-NumpyArray1dT = NDArray[Any]
-RadomStateT = int | np.random.RandomState
-
-GenericParamsT = TypeVar("GenericParamsT", bound=SamplerParamsT)
-
-###########################################################
-# Processing Mixins
-# -----------------
-# These mixins are used to apply pre and post processing
-# to the data source output. At the moment, the only
-# processing that is applied is an optional type/format
-# conversion of the output data.
-#
-# Preprocessing: Applied *before* constraints are applied.
-# Postprocessing: Applied at the end of dataset generation.
-#
-# IMPORTANT: These are only applied when the data are
-# being injected into a DataFrame by the DatasetGenerator.
-###########################################################
-
-
-class PassthroughMixin:
-    @staticmethod
-    def preproc(series: pd.Series, convert_to: str) -> pd.Series:
-        return series
-
-    @staticmethod
-    def postproc(series: pd.Series, convert_to: str) -> pd.Series:
-        return series
-
-    @staticmethod
-    def validate_data_conversion(convert_to: str | None) -> None:
-        pass
-
-
-class TypeConversionMixin:
-    """Converts the data type of the output data.
-
-    This mixin applies the same conversion to both the pre and post
-    processing steps. The preprocessing is needed to ensure constraints
-    are applied to the correct data type. The postprocessing is needed
-    to ensure the final dtype is correct. For example, if the user wants an
-    `int`, we need to convert to `int` before applying constraints, but
-    the ints will be converted back to floats when injected into the
-    dataframe (assuming some rows are non-int values). We therefore need
-    to convert back to `int` after all constraints have been applied.
-    """
-
-    @staticmethod
-    def preproc(series: pd.Series, convert_to: str) -> pd.Series:
-        if convert_to is not None:
-            if convert_to == "int":
-                series = series.round()
-            return series.astype(convert_to)
-        return series
-
-    @staticmethod
-    def postproc(series: pd.Series, convert_to: str | None) -> pd.Series:
-        if convert_to is not None:
-            if convert_to == "int":
-                series = series.round()
-            return series.astype(convert_to)
-        return series
-
-    @staticmethod
-    def validate_data_conversion(convert_to: str | None) -> None:
-        if convert_to is not None and convert_to not in ["float", "int", "str"]:
-            raise ValueError(f"Invalid `convert_to` value: {convert_to}. Must be one of: [float, int, str]")
-
-
-class DatetimeFormatMixin:
-    @staticmethod
-    def preproc(series: pd.Series, convert_to: str | None) -> pd.Series:
-        return series
-
-    @staticmethod
-    def postproc(series: pd.Series, convert_to: str | None) -> pd.Series:
-        if convert_to is not None:
-            return series.dt.strftime(convert_to)
-        if series.dt.month.nunique() == 1:
-            return series.apply(lambda dt: dt.year).astype(str)
-        if series.dt.day.nunique() == 1:
-            return series.apply(lambda dt: dt.strftime("%Y-%m"))
-        if series.dt.hour.sum() > 0 or series.dt.minute.sum() > 0:
-            return series.apply(lambda dt: dt.isoformat()).astype(str)
-        if series.dt.second.sum() == 0:
-            return series.apply(lambda dt: dt.date()).astype(str)
-        return series.apply(lambda dt: dt.isoformat()).astype(str)
-
-    @staticmethod
-    def validate_data_conversion(convert_to: str | None) -> None:
-        if convert_to is not None:
-            try:
-                pd.to_datetime(pd.to_datetime("2012-12-21").strftime(convert_to))
-            except Exception as e:
-                raise ValueError(f"Invalid datetime format: {convert_to}. {e}")
-
-
-###########################################################
-# Base Data Source Classes
-###########################################################
-
-
-class DataSource(ABC, Generic[GenericParamsT]):
-    def __init__(
-        self,
-        params: GenericParamsT,
-        random_state: RadomStateT | None = None,
-        **kwargs,
-    ):
-        self.rng = check_random_state(random_state)
-        self.params = self.get_param_type().model_validate(params)
-        self._setup(**kwargs)
-        self._validate()
-
-    @classmethod
-    def get_param_type(cls) -> type[GenericParamsT]:
-        return cls.__orig_bases__[-1].__args__[0]
-
-    @abstractmethod
-    def inject_data_column(
-        self,
-        dataframe: pd.DataFrame,
-        column_name: str,
-        index: list[int] | None = None,
-    ) -> pd.DataFrame: ...
-
-    @staticmethod
-    @abstractmethod
-    def preproc(series: pd.Series) -> pd.Series: ...
-
-    @staticmethod
-    @abstractmethod
-    def postproc(series: pd.Series, convert_to: str | None) -> pd.Series: ...
-
-    @staticmethod
-    @abstractmethod
-    def validate_data_conversion(convert_to: str | None) -> None: ...
-
-    def get_required_column_names(self) -> tuple[str, ...]:
-        return tuple()
-
-    def _setup(self, **kwargs) -> None:
-        pass
-
-    def _validate(self) -> None:
-        pass
-
-
-class Sampler(DataSource[GenericParamsT], ABC):
-    def _recast_types_if_needed(
-        self,
-        index: list[int] | slice,
-        column_name: str,
-        sample: NumpyArray1dT,
-        dataframe: pd.DataFrame,
-    ) -> pd.DataFrame:
-        # Type may be different if the column has mixed types / NaNs.
-        if column_name in dataframe.columns:
-            dtype = sample.dtype
-            if dtype != dataframe.loc[index, column_name].dtype:
-                dataframe = dataframe.astype({column_name: dtype}, errors="ignore")
-        return dataframe
-
-    def inject_data_column(
-        self,
-        dataframe: pd.DataFrame,
-        column_name: str,
-        index: list[int] | None = None,
-    ) -> pd.DataFrame:
-        index = slice(None) if index is None else index
-
-        if len(index) == 0:
-            return dataframe
-
-        sample = self.sample(len(index))
-
-        # Try recasting before assigning the sample to the dataframe, since setting an item
-        # of incompatible dtype is deprecated and will raise an error in future versions.
-        dataframe = self._recast_types_if_needed(index, column_name, sample, dataframe)
-        dataframe.loc[index, column_name] = sample
-
-        # Recast again in case the assignment led to inconsistencies (e.g., funny business from NaNs).
-        dataframe = self._recast_types_if_needed(index, column_name, sample, dataframe)
-
-        return dataframe
-
-    @abstractmethod
-    def sample(self, num_samples: int) -> NumpyArray1dT: ...
-
-
-class ScipyStatsSampler(Sampler[GenericParamsT], ABC):
-    @property
-    @abstractmethod
-    def distribution(self) -> scipy.stats.rv_continuous | scipy.stats.rv_discrete: ...
-
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        return self.distribution.rvs(size=num_samples, random_state=self.rng)

data_designer/engine/sampling_gen/data_sources/errors.py

@@ -1,12 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from data_designer.engine.sampling_gen.errors import SamplingGenError
-
-
-class InvalidSamplerParamsError(SamplingGenError): ...
-
-
-class PersonSamplerConstraintsError(SamplingGenError): ...

data_designer/engine/sampling_gen/data_sources/sources.py

@@ -1,347 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import uuid
-from typing import TYPE_CHECKING
-
-from data_designer.config.sampler_params import (
-    BernoulliMixtureSamplerParams,
-    BernoulliSamplerParams,
-    BinomialSamplerParams,
-    CategorySamplerParams,
-    DatetimeSamplerParams,
-    GaussianSamplerParams,
-    PersonFromFakerSamplerParams,
-    PersonSamplerParams,
-    PoissonSamplerParams,
-    SamplerParamsT,
-    SamplerType,
-    ScipySamplerParams,
-    SubcategorySamplerParams,
-    TimeDeltaSamplerParams,
-    UniformSamplerParams,
-    UUIDSamplerParams,
-)
-from data_designer.engine.sampling_gen.data_sources.base import (
-    DataSource,
-    DatetimeFormatMixin,
-    NumpyArray1dT,
-    PassthroughMixin,
-    Sampler,
-    ScipyStatsSampler,
-    TypeConversionMixin,
-)
-from data_designer.engine.sampling_gen.data_sources.errors import (
-    InvalidSamplerParamsError,
-    PersonSamplerConstraintsError,
-)
-from data_designer.engine.sampling_gen.entities.dataset_based_person_fields import PERSONA_FIELDS, PII_FIELDS
-from data_designer.engine.sampling_gen.people_gen import PeopleGen
-from data_designer.lazy_heavy_imports import np, pd, scipy
-
-if TYPE_CHECKING:
-    import numpy as np
-    import pandas as pd
-    import scipy
-
-ONE_BILLION = 10**9
-
-
-class SamplerRegistry:
-    _registry: dict[str, type] = {}
-    _reverse_registry: dict[type, str] = {}
-    _params_registry: dict[type, type] = {}
-
-    @classmethod
-    def register(cls, alias: str):
-        def decorator(wrapped_class: type[DataSource[SamplerParamsT]]) -> type:
-            cls._registry[alias] = wrapped_class
-            cls._reverse_registry[wrapped_class] = alias
-            cls._params_registry[wrapped_class.get_param_type()] = wrapped_class
-            return wrapped_class
-
-        return decorator
-
-    @classmethod
-    def get_sampler(cls, alias: str) -> type[DataSource[SamplerParamsT]]:
-        return cls._registry[alias.lower()]
-
-    @classmethod
-    def get_sampler_for_params(cls, params_type: SamplerParamsT) -> type[DataSource[SamplerParamsT]]:
-        return cls._params_registry[type(params_type)]
-
-    @classmethod
-    def get_sampler_alias_for_params(cls, params_type: SamplerParamsT) -> str:
-        return cls._reverse_registry[cls._params_registry[type(params_type)]]
-
-    @classmethod
-    def is_registered(cls, alias: str) -> bool:
-        return alias in cls._registry
-
-    @classmethod
-    def validate_sampler_type(
-        cls, sampler_type: str | type[DataSource[SamplerParamsT]]
-    ) -> type[DataSource[SamplerParamsT]]:
-        if isinstance(sampler_type, str):
-            if sampler_type not in cls._registry:
-                raise ValueError(
-                    f"Sampler type `{sampler_type}` not found in the registry. "
-                    f"Available samplers: {list(cls._registry.keys())}"
-                )
-            sampler_type = cls.get_sampler(sampler_type)
-        if not issubclass(sampler_type, DataSource):
-            raise ValueError(f"Sampler type `{sampler_type}` is not a subclass of `DataSource`")
-        return sampler_type
-
-
-#########################################
-# Data Source Subclasses
-#########################################
-
-
-@SamplerRegistry.register(SamplerType.SUBCATEGORY)
-class SubcategorySampler(TypeConversionMixin, DataSource[SubcategorySamplerParams]):
-    def get_required_column_names(self) -> tuple[str, ...]:
-        return (self.params.category,)
-
-    def inject_data_column(
-        self,
-        dataframe: pd.DataFrame,
-        column_name: str,
-        index: list[int] | None = None,
-    ) -> pd.DataFrame:
-        index = slice(None) if index is None else index
-
-        if len(index) == 0:
-            return dataframe
-
-        dataframe.loc[index, column_name] = dataframe.loc[index, self.params.category].apply(
-            lambda cat_value: self.rng.choice(self.params.values[cat_value])
-        )
-
-        return dataframe
-
-
-#########################################
-# Sampler Subclasses
-#########################################
-
-
-@SamplerRegistry.register(SamplerType.CATEGORY)
-class CategorySampler(TypeConversionMixin, Sampler[CategorySamplerParams]):
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        return self.rng.choice(self.params.values, size=num_samples, p=self.params.weights)
-
-
-@SamplerRegistry.register(SamplerType.DATETIME)
-class DatetimeSampler(DatetimeFormatMixin, Sampler[DatetimeSamplerParams]):
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        # Convert nanoseconds to seconds.
-        start_sec = pd.to_datetime(self.params.start).value // ONE_BILLION
-        end_sec = pd.to_datetime(self.params.end).value // ONE_BILLION
-
-        random_ns = (ONE_BILLION * self.rng.randint(start_sec, end_sec, num_samples, dtype=np.int64)).view(
-            "datetime64[ns]"
-        )
-
-        return np.array(random_ns, dtype=f"datetime64[{self.params.unit}]")
-
-
-@SamplerRegistry.register(SamplerType.PERSON)
-class PersonSampler(PassthroughMixin, Sampler[PersonSamplerParams]):
-    def _setup(self, **kwargs) -> None:
-        self._generator = None
-        self._fixed_kwargs = {}
-        for field in self.params.generator_kwargs:
-            if getattr(self.params, field) is not None:
-                attr = getattr(self.params, field)
-                if field == "select_field_values":
-                    for key, value in attr.items():
-                        if key == "state" and self.params.locale == "en_US":
-                            key = "region"  # This is the field name in the census-based person dataset.
-                        if key not in PII_FIELDS + PERSONA_FIELDS:
-                            raise ValueError(f"Invalid field name: {key}")
-                        self._fixed_kwargs[key] = value
-                else:
-                    self._fixed_kwargs[field] = attr
-        if people_gen_resource := kwargs.get("people_gen_resource"):
-            if self.params.people_gen_key not in people_gen_resource:
-                raise ValueError(f"Person generator with key {self.params.people_gen_key} not found.")
-            self.set_generator(people_gen_resource[self.params.people_gen_key])
-
-    def set_generator(self, generator: PeopleGen) -> None:
-        self._generator = generator
-
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        if self._generator is None:
-            raise ValueError("Generator not set. Please setup generator before sampling.")
-
-        samples = np.array(self._generator.generate(num_samples, **self._fixed_kwargs))
-        if len(samples) < num_samples:
-            raise PersonSamplerConstraintsError(
-                f"🛑 Only {len(samples)} samples could be generated with the given settings: {self._fixed_kwargs!r}. "
-                "This is likely because the filter values are too strict. Person sampling does not support "
-                "rare combinations of field values. Please loosen the constraints and try again."
-            )
-        return samples
-
-
-@SamplerRegistry.register(SamplerType.PERSON_FROM_FAKER)
-class PersonFromFakerSampler(PassthroughMixin, Sampler[PersonFromFakerSamplerParams]):
-    def _setup(self, **kwargs) -> None:
-        self._generator = None
-        self._fixed_kwargs = {}
-        for field in self.params.generator_kwargs:
-            if getattr(self.params, field) is not None:
-                self._fixed_kwargs[field] = getattr(self.params, field)
-        if people_gen_resource := kwargs.get("people_gen_resource"):
-            if self.params.people_gen_key not in people_gen_resource:
-                raise ValueError(f"Person generator with key {self.params.people_gen_key} not found.")
-            self.set_generator(people_gen_resource[self.params.people_gen_key])
-
-    def set_generator(self, generator: PeopleGen) -> None:
-        self._generator = generator
-
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        if self._generator is None:
-            raise ValueError("Generator not set. Please setup generator before sampling.")
-
-        samples = np.array(self._generator.generate(num_samples, **self._fixed_kwargs))
-        if len(samples) < num_samples:
-            raise ValueError(f"Only {len(samples)} samples could be generated given constraints {self._fixed_kwargs}.")
-        return samples
-
-
-@SamplerRegistry.register(SamplerType.TIMEDELTA)
-class TimeDeltaSampler(DatetimeFormatMixin, Sampler[TimeDeltaSamplerParams]):
-    def get_required_column_names(self) -> tuple[str, ...]:
-        return (self.params.reference_column_name,)
-
-    def inject_data_column(
-        self,
-        dataframe: pd.DataFrame,
-        column_name: str,
-        index: list[int] | None = None,
-    ) -> pd.DataFrame:
-        index = slice(None) if index is None else index
-
-        if self.params.reference_column_name not in list(dataframe):
-            raise ValueError(f"Columns `{self.params.reference_column_name}` not found in dataset")
-
-        dataframe.loc[index, column_name] = pd.to_datetime(
-            dataframe.loc[index, self.params.reference_column_name]
-        ) + pd.to_timedelta(self.sample(len(index)), unit=self.params.unit)
-
-        return dataframe
-
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        deltas = self.rng.randint(self.params.dt_min, self.params.dt_max, num_samples)
-        return np.array(deltas, dtype=f"timedelta64[{self.params.unit}]")
-
-
-@SamplerRegistry.register(SamplerType.UUID)
-class UUIDSampler(PassthroughMixin, Sampler[UUIDSamplerParams]):
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        prefix = self.params.prefix or ""
-
-        uid_list = []
-        while len(uid_list) < num_samples:
-            uid = (
-                f"{prefix}{uuid.uuid4().hex[: self.params.last_index].upper()}"
-                if self.params.uppercase
-                else f"{prefix}{uuid.uuid4().hex[: self.params.last_index]}"
-            )
-            if uid not in uid_list:
-                uid_list.append(uid)
-
-        return np.array(uid_list)
-
-
-#########################################
-# Scipy Samplers
-#########################################
-
-
-@SamplerRegistry.register(SamplerType.SCIPY)
-class ScipySampler(TypeConversionMixin, ScipyStatsSampler[ScipySamplerParams]):
-    """Escape hatch sampler to give users access to any scipy.stats distribution."""
-
-    @property
-    def distribution(self) -> scipy.stats.rv_continuous | scipy.stats.rv_discrete:
-        return getattr(scipy.stats, self.params.dist_name)(**self.params.dist_params)
-
-    def _validate(self) -> None:
-        _validate_scipy_distribution(self.params.dist_name, self.params.dist_params)
-
-
-@SamplerRegistry.register(SamplerType.BERNOULLI)
-class BernoulliSampler(TypeConversionMixin, ScipyStatsSampler[BernoulliSamplerParams]):
-    @property
-    def distribution(self) -> scipy.stats.rv_discrete:
-        return scipy.stats.bernoulli(p=self.params.p)
-
-
-@SamplerRegistry.register(SamplerType.BERNOULLI_MIXTURE)
-class BernoulliMixtureSampler(TypeConversionMixin, Sampler[BernoulliMixtureSamplerParams]):
-    def sample(self, num_samples: int) -> NumpyArray1dT:
-        return scipy.stats.bernoulli(p=self.params.p).rvs(size=num_samples) * getattr(
-            scipy.stats, self.params.dist_name
-        )(**self.params.dist_params).rvs(size=num_samples)
-
-    def _validate(self) -> None:
-        _validate_scipy_distribution(self.params.dist_name, self.params.dist_params)
-
-
-@SamplerRegistry.register(SamplerType.BINOMIAL)
-class BinomialSampler(TypeConversionMixin, ScipyStatsSampler[BinomialSamplerParams]):
-    @property
-    def distribution(self) -> scipy.stats.rv_discrete:
-        return scipy.stats.binom(n=self.params.n, p=self.params.p)
-
-
-@SamplerRegistry.register(SamplerType.GAUSSIAN)
-class GaussianSampler(TypeConversionMixin, ScipyStatsSampler[GaussianSamplerParams]):
-    @property
-    def distribution(self) -> scipy.stats.rv_continuous:
-        return scipy.stats.norm(loc=self.params.mean, scale=self.params.stddev)
-
-
-@SamplerRegistry.register(SamplerType.POISSON)
-class PoissonSampler(TypeConversionMixin, ScipyStatsSampler[PoissonSamplerParams]):
-    @property
-    def distribution(self) -> scipy.stats.rv_discrete:
-        return scipy.stats.poisson(mu=self.params.mean)
-
-
-@SamplerRegistry.register(SamplerType.UNIFORM)
-class UniformSampler(TypeConversionMixin, ScipyStatsSampler[UniformSamplerParams]):
-    @property
-    def distribution(self) -> scipy.stats.rv_continuous:
-        return scipy.stats.uniform(loc=self.params.low, scale=self.params.high - self.params.low)
-
-
-###################################################
-# Helper functions for loading sources in isolation
-###################################################
-
-
-def load_sampler(sampler_type: SamplerType, **params) -> DataSource:
-    """Load a data source from a source type and parameters."""
-    return SamplerRegistry.validate_sampler_type(sampler_type)(params=params)
-
-
-def _validate_scipy_distribution(dist_name: str, dist_params: dict) -> None:
-    if not hasattr(scipy.stats, dist_name):
-        raise InvalidSamplerParamsError(f"Distribution {dist_name} not found in scipy.stats")
-    if not hasattr(getattr(scipy.stats, dist_name), "rvs"):
-        raise InvalidSamplerParamsError(
-            f"Distribution {dist_name} does not have a `rvs` method, which is required for sampling."
-        )
-    try:
-        getattr(scipy.stats, dist_name)(**dist_params)
-    except Exception:
-        raise InvalidSamplerParamsError(
-            f"Distribution parameters {dist_params} are not a valid for distribution '{dist_name}'"
-        )

data_designer/engine/sampling_gen/entities/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0