humalab 0.1.0__py3-none-any.whl

humalab/dists/categorical.py

@@ -0,0 +1,81 @@
+from humalab.dists.distribution import Distribution
+
+import numpy as np
+
+class Categorical(Distribution):
+    """Categorical distribution for discrete choices.
+
+    Samples from a list of choices with optional weights. If weights are not
+    provided, samples uniformly from all choices. Weights are automatically
+    normalized to sum to 1.
+    """
+    def __init__(self,
+                 generator: np.random.Generator,
+                 choices: list, 
+                 weights: list[float] | None = None,
+                 size: int | tuple[int, ...] | None = None) -> None:
+        """
+        Initialize the categorical distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            choices (list): The list of choices.
+            weights (list[float] | None): The weights for each choice.
+            size (int | tuple[int, ...] | None): The size of the output.
+        """
+        super().__init__(generator=generator)
+        self._choices = choices
+        self._size = size
+        if weights is not None and not np.isclose(sum(weights), 1.0):
+            weight_sum = sum(weights)
+            weights = [w / weight_sum for w in weights]
+        self._weights = weights
+
+    @staticmethod
+    def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (choices, weights).
+
+        Returns:
+            bool: Always returns True as categorical accepts any parameters.
+        """
+        return True
+
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the categorical distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled choice(s) from the list.
+        """
+        return self._generator.choice(self._choices, size=self._size, p=self._weights)
+
+    def __repr__(self) -> str:
+        """String representation of the categorical distribution.
+
+        Returns:
+            str: String representation showing choices, size, and weights.
+        """
+        return f"Categorical(choices={self._choices}, size={self._size}, weights={self._weights})"
+    
+    @staticmethod
+    def create(generator: np.random.Generator, 
+               choices: list, 
+               weights: list[float] | None = None,
+               size: int | tuple[int, ...] | None = None
+               ) -> 'Categorical':
+        """
+        Create a categorical distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            choices (list): The list of choices.
+            size (int | tuple[int, ...] | None): The size of the output.
+            weights (list[float] | None): The weights for each choice.
+
+        Returns:
+            Categorical: The created categorical distribution.
+        """
+        return Categorical(generator=generator, choices=choices, size=size, weights=weights)

humalab/dists/discrete.py

@@ -0,0 +1,103 @@
+from humalab.dists.distribution import Distribution
+from typing import Any
+
+import numpy as np
+
+class Discrete(Distribution):
+    """Discrete uniform distribution over integers.
+
+    Samples integer values uniformly from a range [low, high). The endpoint
+    parameter controls whether the upper bound is inclusive or exclusive.
+    Supports scalar outputs as well as multi-dimensional arrays with 1D variants.
+    """
+    def __init__(self, 
+                 generator: np.random.Generator,
+                 low: int | Any, 
+                 high: int | Any,
+                 endpoint: bool | None = None,
+                 size: int | tuple[int, ...] | None = None,
+                 ) -> None:
+        """
+        Initialize the discrete distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            low (int | Any): The lower bound (inclusive).
+            high (int | Any): The upper bound (exclusive).
+            endpoint (bool | None): Whether to include the endpoint.
+            size (int | tuple[int, ...] | None): The size of the output.
+        """
+        super().__init__(generator=generator)
+        self._low = np.array(low)
+        self._high = np.array(high)
+        self._size = size
+        self._endpoint = endpoint if endpoint is not None else True
+    
+    @staticmethod
+    def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
+        arg1 = args[0]
+        arg2 = args[1]
+        if dimensions == 0:
+            if not isinstance(arg1, int):
+                return False
+            if not isinstance(arg2, int):
+                return False
+            return True
+        if dimensions == -1:
+            return True
+        if not isinstance(arg1, int):
+            if isinstance(arg1, (list, np.ndarray)):
+                if len(arg1) != dimensions:
+                    return False
+        if not isinstance(arg2, int):
+            if isinstance(arg2, (list, np.ndarray)):
+                if len(arg2) != dimensions:
+                    return False
+        return True
+
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the discrete distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled integer value(s) from [low, high).
+        """
+        return self._generator.integers(self._low, self._high, size=self._size, endpoint=self._endpoint)
+
+    def __repr__(self) -> str:
+        """String representation of the discrete distribution.
+
+        Returns:
+            str: String representation showing low, high, size, and endpoint.
+        """
+        return f"Discrete(low={self._low}, high={self._high}, size={self._size}, endpoint={self._endpoint})"
+    
+    @staticmethod
+    def create(generator: np.random.Generator, 
+               low: int | Any, 
+               high: int | Any, 
+               endpoint: bool = True,
+               size: int | tuple[int, ...] | None = None, 
+               ) -> 'Discrete':
+        """
+        Create a discrete distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            low (int | Any): The lower bound (inclusive).
+            high (int | Any): The upper bound (exclusive).
+            endpoint (bool): Whether to include the endpoint.
+            size (int | tuple[int, ...] | None): The size of the output.
+
+        Returns:
+            Discrete: The created discrete distribution.
+        """
+        return Discrete(generator=generator, low=low, high=high, size=size, endpoint=endpoint)

humalab/dists/distribution.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+class Distribution(ABC):
+    """Abstract base class for probability distributions.
+
+    All distribution classes inherit from this base class and must implement
+    the _sample() method. Distributions maintain a random number generator
+    and track the last sampled value.
+    """
+    def __init__(self,
+                 generator: np.random.Generator) -> None:
+        """
+        Initialize the distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+        """
+        super().__init__()
+        self._generator = generator
+        self._last_sample = None
+
+    def sample(self) -> int | float | np.ndarray:
+        """
+        Sample from the distribution.
+
+        Returns:
+            int | float | np.ndarray: The sampled value(s).
+        """
+        self._last_sample = self._sample()
+        return self._last_sample
+
+    @abstractmethod
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the distribution.
+
+        Returns:
+            int | float | np.ndarray: The sampled value(s).
+        """
+        pass
+
+    @property
+    def last_sample(self) -> int | float | np.ndarray | None:
+        """Get the last sampled value.
+        Returns:
+            int | float | np.ndarray | None: The last sampled value, or None if no sample has been taken yet.
+        """
+        return self._last_sample

humalab/dists/gaussian.py

@@ -0,0 +1,96 @@
+from humalab.dists.distribution import Distribution
+from typing import Any
+import numpy as np
+
+
+class Gaussian(Distribution):
+    """Gaussian (normal) distribution.
+
+    Samples values from a normal distribution with specified mean (loc) and
+    standard deviation (scale). Supports scalar outputs as well as multi-dimensional
+    arrays with 1D, 2D, or 3D variants.
+    """
+    def __init__(self,
+                 generator: np.random.Generator,
+                 loc: float | Any,
+                 scale: float | Any,
+                 size: int | tuple[int, ...] | None = None) -> None:
+        """
+        Initialize the Gaussian (normal) distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            loc (float | Any): The mean of the distribution.
+            scale (float | Any): The standard deviation of the distribution.
+            size (int | tuple[int, ...] | None): The size of the output.
+        """
+        super().__init__(generator=generator)
+        self._loc = loc
+        self._scale = scale
+        self._size = size
+
+    @staticmethod
+    def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (loc, scale).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
+        arg1 = args[0]
+        arg2 = args[1]
+        if dimensions == 0:
+            if not isinstance(arg1, (int, float)):
+                return False
+            if not isinstance(arg2, (int, float)):
+                return False
+            return True
+        if dimensions == -1:
+            return True
+        if not isinstance(arg1, (int, float)):
+            if isinstance(arg1, (list, np.ndarray)):
+                if len(arg1) != dimensions:
+                    return False
+        if not isinstance(arg2, (int, float)):
+            if isinstance(arg2, (list, np.ndarray)):
+                if len(arg2) != dimensions:
+                    return False
+        return True
+
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the Gaussian distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) from N(loc, scale).
+        """
+        return self._generator.normal(loc=self._loc, scale=self._scale, size=self._size)
+
+    def __repr__(self) -> str:
+        """String representation of the Gaussian distribution.
+
+        Returns:
+            str: String representation showing loc, scale, and size.
+        """
+        return f"Gaussian(loc={self._loc}, scale={self._scale}, size={self._size})"
+    
+    @staticmethod
+    def create(generator: np.random.Generator, 
+               loc: float | Any, 
+               scale: float | Any, 
+               size: int | tuple[int, ...] | None = None) -> 'Gaussian':
+        """
+        Create a Gaussian (normal) distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            loc (float | Any): The mean of the distribution.
+            scale (float | Any): The standard deviation of the distribution.
+            size (int | tuple[int, ...] | None): The size of the output.
+
+        Returns:
+            Gaussian: The created Gaussian distribution.
+        """
+        return Gaussian(generator=generator, loc=loc, scale=scale, size=size)

humalab/dists/log_uniform.py

@@ -0,0 +1,97 @@
+from humalab.dists.distribution import Distribution
+from typing import Any
+
+import numpy as np
+
+class LogUniform(Distribution):
+    """Log-uniform distribution.
+
+    Samples values uniformly in log-space, useful for hyperparameters that
+    span multiple orders of magnitude (e.g., learning rates). The result is
+    exp(uniform(log(low), log(high))). Supports scalar outputs as well as
+    multi-dimensional arrays with 1D variants.
+    """
+    def __init__(self, 
+                 generator: np.random.Generator,
+                 low: float | Any,
+                 high: float | Any,
+                 size: int | tuple[int, ...]| None = None) -> None:
+        """
+        Initialize the log-uniform distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            low (float | Any): The lower bound (inclusive).
+            high (float | Any): The upper bound (exclusive).
+            size (int | tuple[int, ...]| None): The size of the output.
+        """
+        super().__init__(generator=generator)
+        self._log_low = np.log(np.array(low))
+        self._log_high = np.log(np.array(high))
+        self._size = size
+    
+    @staticmethod
+    def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
+        arg1 = args[0]
+        arg2 = args[1]
+        if dimensions == 0:
+            if not isinstance(arg1, (int, float)):
+                return False
+            if not isinstance(arg2, (int, float)):
+                return False
+            return True
+        if dimensions == -1:
+            return True
+        if not isinstance(arg1, (int, float)):
+            if isinstance(arg1, (list, np.ndarray)):
+                if len(arg1) != dimensions:
+                    return False
+        if not isinstance(arg2, (int, float)):
+            if isinstance(arg2, (list, np.ndarray)):
+                if len(arg2) != dimensions:
+                    return False
+        return True
+
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the log-uniform distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) in log-space.
+        """
+        return np.exp(self._generator.uniform(self._log_low, self._log_high, size=self._size))
+
+    def __repr__(self) -> str:
+        """String representation of the log-uniform distribution.
+
+        Returns:
+            str: String representation showing low, high, and size.
+        """
+        return f"LogUniform(low={np.exp(self._log_low)}, high={np.exp(self._log_high)}, size={self._size})"
+
+    @staticmethod
+    def create(generator: np.random.Generator, 
+               low: float | Any, 
+               high: float | Any, 
+               size: int | tuple[int, ...]| None = None) -> 'LogUniform':
+        """
+        Create a log-uniform distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            low (float | Any): The lower bound (inclusive).
+            high (float | Any): The upper bound (exclusive).
+            size (int | tuple[int, ...]| None): The size of the output.
+
+        Returns:
+            LogUniform: The created log-uniform distribution.
+        """
+        return LogUniform(generator=generator, low=low, high=high, size=size)

humalab/dists/truncated_gaussian.py

@@ -0,0 +1,129 @@
+from humalab.dists.distribution import Distribution
+from typing import Any
+import numpy as np
+
+
+class TruncatedGaussian(Distribution):
+    """Truncated Gaussian (normal) distribution.
+
+    Samples values from a normal distribution with specified mean (loc) and
+    standard deviation (scale), but constrained to lie within [low, high].
+    Values outside the bounds are resampled until they fall within range.
+    Supports scalar outputs as well as multi-dimensional arrays with 1D, 2D, or 3D variants.
+    """
+    def __init__(self,
+                 generator: np.random.Generator,
+                 loc: float | Any,
+                 scale: float | Any,
+                 low: float | Any,
+                 high: float | Any,
+                 size: int | tuple[int, ...] | None = None) -> None:
+        """
+        Initialize the truncated Gaussian (normal) distribution.
+        
+        Args:
+            generator (np.random.Generator): The random number generator.
+            loc (float | Any): The mean of the distribution.
+            scale (float | Any): The standard deviation of the distribution.
+            low (float | Any): The lower truncation bound.
+            high (float | Any): The upper truncation bound.
+            size (int | tuple[int, ...] | None): The size of the output.
+        """
+        super().__init__(generator=generator)
+        self._loc = loc
+        self._scale = scale
+        self._low = low
+        self._high = high
+        self._size = size
+
+    @staticmethod
+    def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (loc, scale, low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
+        arg1 = args[0]
+        arg2 = args[1]
+        arg3 = args[2]
+        arg4 = args[3]
+        if dimensions == 0:
+            if not isinstance(arg1, (int, float)):
+                return False
+            if not isinstance(arg2, (int, float)):
+                return False
+            if not isinstance(arg3, (int, float)):
+                return False
+            if not isinstance(arg4, (int, float)):
+                return False
+            return True
+        if dimensions == -1:
+            return True
+        if not isinstance(arg1, (int, float)):
+            if isinstance(arg1, (list, np.ndarray)):
+                if len(arg1) != dimensions:
+                    return False
+        if not isinstance(arg2, (int, float)):
+            if isinstance(arg2, (list, np.ndarray)):
+                if len(arg2) != dimensions:
+                    return False
+        if not isinstance(arg3, (int, float)):
+            if isinstance(arg3, (list, np.ndarray)):
+                if len(arg3) != dimensions:
+                    return False
+        if not isinstance(arg4, (int, float)):
+            if isinstance(arg4, (list, np.ndarray)):
+                if len(arg4) != dimensions:
+                    return False
+        return True
+    
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the truncated Gaussian distribution.
+
+        Samples are generated from N(loc, scale) and resampled if they fall
+        outside [low, high].
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) within [low, high].
+        """
+        samples = self._generator.normal(loc=self._loc, scale=self._scale, size=self._size)
+        mask = (samples < self._low) | (samples > self._high)
+        while np.any(mask):
+            samples[mask] = self._generator.normal(loc=self._loc, scale=self._scale, size=np.sum(mask))
+            mask = (samples < self._low) | (samples > self._high)
+        return samples
+
+    def __repr__(self) -> str:
+        """String representation of the truncated Gaussian distribution.
+
+        Returns:
+            str: String representation showing loc, scale, low, high, and size.
+        """
+        return f"TruncatedGaussian(loc={self._loc}, scale={self._scale}, low={self._low}, high={self._high}, size={self._size})"
+    
+    @staticmethod
+    def create(generator: np.random.Generator, 
+               loc: float | Any, 
+               scale: float | Any, 
+               low: float | Any, 
+               high: float | Any, 
+               size: int | tuple[int, ...] | None = None) -> 'TruncatedGaussian':
+        """
+        Create a truncated Gaussian (normal) distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            loc (float | Any): The mean of the distribution.
+            scale (float | Any): The standard deviation of the distribution.
+            low (float | Any): The lower truncation bound.
+            high (float | Any): The upper truncation bound.
+            size (int | tuple[int, ...] | None): The size of the output.
+
+        Returns:
+            TruncatedGaussian: The created truncated Gaussian distribution.
+        """
+        return TruncatedGaussian(generator=generator, loc=loc, scale=scale, low=low, high=high, size=size)

humalab/dists/uniform.py

@@ -0,0 +1,95 @@
+from humalab.dists.distribution import Distribution
+
+from typing import Any
+import numpy as np
+
+class Uniform(Distribution):
+    """Uniform distribution over a continuous or discrete range.
+
+    Samples values uniformly from the half-open interval [low, high). Supports
+    scalar outputs as well as multi-dimensional arrays with 1D, 2D, or 3D variants.
+    """
+    def __init__(self, 
+                 generator: np.random.Generator,
+                 low: float | Any, 
+                 high: float | Any, 
+                 size: int | tuple[int, ...] | None = None, ) -> None:
+        """
+        Initialize the uniform distribution.
+        
+        Args:
+            generator (np.random.Generator): The random number generator.
+            low (float | Any): The lower bound (inclusive).
+            high (float | Any): The upper bound (exclusive).
+            size (int | tuple[int, ...] | None): The size of the output.
+        """
+        super().__init__(generator=generator)
+        self._low = np.array(low)
+        self._high = np.array(high)
+        self._size = size
+
+    @staticmethod
+    def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
+        arg1 = args[0]
+        arg2 = args[1]
+        if dimensions == 0:
+            if not isinstance(arg1, (int, float)):
+                return False
+            if not isinstance(arg2, (int, float)):
+                return False
+            return True
+        if dimensions == -1:
+            return True
+        if not isinstance(arg1, (int, float)):
+            if isinstance(arg1, (list, np.ndarray)):
+                if len(arg1) > dimensions:
+                    return False
+        if not isinstance(arg2, (int, float)):
+            if isinstance(arg2, (list, np.ndarray)):
+                if len(arg2) > dimensions:
+                    return False
+        return True
+
+    def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the uniform distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) from [low, high).
+        """
+        return self._generator.uniform(self._low, self._high, size=self._size)
+
+    def __repr__(self) -> str:
+        """String representation of the uniform distribution.
+
+        Returns:
+            str: String representation showing low, high, and size.
+        """
+        return f"Uniform(low={self._low}, high={self._high}, size={self._size})"
+    
+    @staticmethod
+    def create(generator: np.random.Generator, 
+               low: float | Any, 
+               high: float | Any, 
+               size: int | tuple[int, ...] | None = None) -> 'Uniform':
+        """
+        Create a uniform distribution.
+
+        Args:
+            generator (np.random.Generator): The random number generator.
+            low (float | Any): The lower bound (inclusive).
+            high (float | Any): The upper bound (exclusive).
+            size (int | tuple[int, ...] | None): The size of the output.
+
+        Returns:
+            Uniform: The created uniform distribution.
+        """
+        return Uniform(generator=generator, low=low, high=high, size=size)