procfunc 0.30.0__py3-none-any.whl

procfunc/random.py

@@ -0,0 +1,247 @@
+import logging
+from typing import TypeVar
+
+import numpy as np
+
+import procfunc as pf
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+@pf.tracer.primitive(allow_exec=True)
+def uniform(
+    rng: np.random.Generator,
+    low: float,
+    high: float,
+) -> float:
+    return rng.uniform(low, high)
+
+
+@pf.tracer.primitive(allow_exec=True)
+def normal(
+    rng: np.random.Generator,
+    mean: float,
+    std: float,
+) -> float:
+    return rng.normal(mean, std)
+
+
+@pf.tracer.primitive(allow_exec=True)
+def randint(
+    rng: np.random.Generator,
+    low: int,
+    high: int,
+) -> int:
+    return rng.integers(low, high)
+
+
+@pf.tracer.primitive(allow_exec=True)
+def uniform_tails(
+    rng: np.random.Generator,
+    low: float,
+    high: float,
+    tail_pct: float = 0.05,
+) -> float:
+    if rng.random() < 0.5:
+        return rng.uniform(low, low + (high - low) * tail_pct)
+    else:
+        return rng.uniform(high - (high - low) * tail_pct, high)
+
+
+@pf.tracer.primitive(allow_exec=True)
+def clip_gaussian(
+    rng: np.random.Generator,
+    mean: float,
+    std: float,
+    low: float | None = None,
+    high: float | None = None,
+    max_tries: int = 20,
+) -> float:
+    low = low or mean - 3 * std
+    high = high or mean + 3 * std
+
+    for _ in range(max_tries):
+        val = rng.normal(mean, std)
+        if low <= val and val <= high:
+            return val
+
+    raise ValueError(
+        f"clip_gaussian({mean=}, {std=}, {low=}, {high=}) failed to "
+        f"sample a value in the range after {max_tries=}"
+    )
+
+
+def wrap_gaussian(
+    rng: np.random.Generator,
+    mean: float,
+    std: float,
+    low: float,
+    high: float,
+):
+    x = rng.normal(mean, std)
+
+    if x < low:
+        x = high - (low - x) % (high - low)
+    elif x > high:
+        x = low + (x - high) % (high - low)
+
+    return x
+
+
+def exponential(
+    rng: np.random.Generator,
+    scale: float,
+) -> float:
+    return rng.exponential(scale)
+
+
+def log_uniform(
+    rng: np.random.Generator,
+    low: float,
+    high: float,
+    size: tuple | None = None,
+) -> float:
+    return np.exp(rng.uniform(np.log(low), np.log(high), size=size))
+
+
+def log_normal(
+    rng: np.random.Generator,
+    mean: float,
+    std: float,
+    size: tuple | None = None,
+) -> float:
+    return np.exp(rng.normal(np.log(mean), std, size=size))
+
+
+def spherical_sample(
+    rng: np.random.Generator,
+    min_elevation: float | None = None,
+    max_elevation: float | None = None,
+    retries: int = 100,
+) -> float:
+    for _ in range(retries):
+        P = rng.standard_normal(3)
+        x = np.arctan2(np.abs(P[2]), (P[0] ** 2 + P[1] ** 2) ** 0.5)
+        if (min_elevation is None or x > np.radians(min_elevation)) and (
+            max_elevation is None or x < np.radians(max_elevation)
+        ):
+            return np.degrees(x)
+
+    raise ValueError(
+        f"spherical_sample({min_elevation=}, {max_elevation=}) failed to "
+        f"sample a value in the range after {retries=}"
+    )
+
+
+def mixture_of_gaussian(
+    rng: np.random.Generator,
+    means: np.ndarray,
+    stds: np.ndarray,
+    weights: list[float],
+    low: np.ndarray | float | None = None,
+    high: np.ndarray | float | None = None,
+):
+    p = np.array(weights) / np.sum(weights)
+    idx = rng.choice(len(p), p=p)
+
+    mu = means[idx]
+    std = stds[idx]
+
+    res = rng.normal(mu, std)
+
+    if low is not None:
+        res = np.maximum(res, low)
+    if high is not None:
+        res = np.minimum(res, high)
+
+    return res
+
+
+def beta(
+    rng: np.random.Generator,
+    a: float,
+    b: float,
+) -> float:
+    return rng.beta(a, b)
+
+
+def poisson(
+    rng: np.random.Generator,
+    lam: float,
+) -> float:
+    return rng.poisson(lam)
+
+
+def triangular(
+    rng: np.random.Generator,
+    low: float,
+    high: float,
+    mode: float,
+) -> float:
+    return rng.triangular(low, mode, high)
+
+
+def gamma(
+    rng: np.random.Generator,
+    shape: float,
+    scale: float,
+) -> float:
+    return rng.gamma(shape, scale)
+
+
+def binomial(
+    rng: np.random.Generator,
+    n: int,
+    p: float,
+) -> float:
+    return rng.binomial(n, p)
+
+
+def geometric(
+    rng: np.random.Generator,
+    p: float,
+) -> float:
+    return rng.geometric(p)
+
+
+random_distrib_funcs = (
+    uniform,
+    normal,
+    uniform_tails,
+    clip_gaussian,
+    wrap_gaussian,
+    log_uniform,
+    log_normal,
+    spherical_sample,
+    mixture_of_gaussian,
+    beta,
+    poisson,
+    triangular,
+    gamma,
+    binomial,
+    geometric,
+    exponential,
+)
+
+
+__all__ = [
+    "uniform",
+    "normal",
+    "uniform_tails",
+    "clip_gaussian",
+    "wrap_gaussian",
+    "log_uniform",
+    "log_normal",
+    "spherical_sample",
+    "mixture_of_gaussian",
+    "beta",
+    "poisson",
+    "triangular",
+    "gamma",
+    "binomial",
+    "geometric",
+    "exponential",
+    "random_distrib_funcs",
+]

procfunc/tracer/__init__.py

@@ -0,0 +1,43 @@
+from procfunc.tracer.decorator import (
+    generator,
+    grammar,
+    primitive,
+    random_control,
+    random_param,
+    register_trace_target,
+)
+from procfunc.tracer.patch import (
+    PATCHING_FLAG_ATTR,
+    Patcher,
+    PatchFunctionTarget,
+    TraceLevel,
+)
+from procfunc.tracer.proxy import (
+    RngProxy,
+    RngSpawnResultProxy,
+)
+from procfunc.tracer.trace import (
+    add_banned_module,
+    add_search_scope,
+    add_wrap_target,
+    autowrap_module,
+    trace,
+)
+
+__all__ = [
+    "TraceLevel",
+    "grammar",
+    "random_control",
+    "random_param",
+    "generator",
+    "primitive",
+    "RngProxy",
+    "RngSpawnResultProxy",
+    "Patcher",
+    "PatchFunctionTarget",
+    "PATCHING_FLAG_ATTR",
+    "trace",
+    "add_banned_module",
+    "add_wrap_target",
+    "autowrap_module",
+]

procfunc/tracer/decorator.py

@@ -0,0 +1,121 @@
+import functools
+import inspect
+import logging
+from typing import Callable
+
+from .patch import PatchFunctionTarget, TraceLevel
+from .trace import add_wrap_target
+
+logger = logging.getLogger(__name__)
+
+
+def register_trace_target(
+    func: Callable,
+    trace_level: TraceLevel,
+    allow_exec: bool,
+    custom_trace_wrapper_create: Callable | None,
+    mutates: list[str] | None = None,
+    normalize: bool = True,
+):
+    """Helper function to register a function for tracing with consistent frame inspection"""
+
+    if ".<locals>." in getattr(func, "__qualname__", ""):
+        raise ValueError(
+            f"Cannot register {func.__qualname__!r} as a trace target: local functions "
+            "(defined inside another function) are not in module globals and cannot be "
+            "traced. Move it to module level or remove the procfunc decorator."
+        )
+
+    f = inspect.currentframe()
+    while f.f_code.co_name != "<module>":
+        assert hasattr(f, "f_back")
+        f = f.f_back
+        assert hasattr(f, "f_code")
+
+    frame = f.f_globals
+    sourcename = f.f_code.co_filename.split("/")[-1][:-3]  # remove .py
+
+    assert hasattr(func, "__name__")
+    assert isinstance(func.__name__, str)
+
+    # Local functions (defined inside another function) are never in the module
+    # globals frame, so registering them would leave a stale entry that causes
+    # KeyError on every subsequent pf.trace() call.
+    if ".<locals>." in getattr(func, "__qualname__", ""):
+        return
+
+    add_wrap_target(
+        PatchFunctionTarget(
+            frame=frame,
+            name=func.__name__,
+            trace_level=trace_level,
+            allow_exec=allow_exec,
+            custom_trace_wrapper_create=custom_trace_wrapper_create,
+            source_name=sourcename,
+            mutates=mutates,
+            normalize=normalize,
+        )
+    )
+
+
+def _make_decorator(
+    trace_level: TraceLevel = TraceLevel.GENERATORS,
+):
+    """
+    Main decorator to mark a function and control its behaviour during tracing
+
+    Args:
+        trace_level: The level at which this function should be traced
+        allow_exec: If True, the function may be executed. Use False for any functions that create meshes or have side-effects
+        custom_trace_wrapper_create: Custom wrapper creation function for special tracing behavior
+    """
+
+    def decorator(
+        func: Callable | None = None,
+        *,
+        allow_exec: bool = False,
+        custom_trace_wrapper_create: Callable | None = None,
+        mutates: list[str] | None = None,
+        normalize: bool = True,
+    ) -> Callable:
+        @functools.wraps(func)
+        def decorate(
+            func: Callable,
+        ):
+            logger.debug(
+                f"Decorating {func.__name__} with {trace_level=} {custom_trace_wrapper_create=}"
+            )
+            # When called as decorator, we need to go back 2 frames
+            register_trace_target(
+                func,
+                trace_level,
+                allow_exec,
+                custom_trace_wrapper_create,
+                mutates=mutates,
+                normalize=normalize,
+            )
+            return func
+
+        # handle EITHER @decorator or @decorator(args)
+        if func is not None:
+            return decorate(func)
+        else:
+            return decorate
+
+    return decorator
+
+
+grammar = _make_decorator(trace_level=TraceLevel.GRAMMAR)
+random_control = _make_decorator(trace_level=TraceLevel.RANDOM_CONTROL)
+random_param = _make_decorator(trace_level=TraceLevel.RANDOM_PARAMS)
+generator = _make_decorator(trace_level=TraceLevel.GENERATORS)
+primitive = _make_decorator(trace_level=TraceLevel.PRIMITIVES)
+
+__all__ = [
+    "register_trace_target",
+    "grammar",
+    "random_control",
+    "random_param",
+    "generator",
+    "primitive",
+]