numba Module

Numba acceleration utilities and compatibility layer.

Overview

The numba subpackage provides njit_switch, RNG functions (binomial, multinomial), and Numba enable/disable controls for the simulation engine.

Complete Module Reference

natal.numba

Numba acceleration subpackage.

Provides configurable JIT compilation decorators (:mod:natal.numba.utils) and Numba-compatible helper functions with dual implementations (:mod:natal.numba.compat).

disable_numba

disable_numba()

Disable Numba JIT compilation globally (use pure Python).

Useful for debugging or when you need Python-compatible error messages. Note: Numba is ENABLED by default.

Source code in src/natal/numba/utils.py

def disable_numba():
    """Disable Numba JIT compilation globally (use pure Python).

    Useful for debugging or when you need Python-compatible error messages.
    Note: Numba is ENABLED by default.
    """
    global NUMBA_ENABLED
    NUMBA_ENABLED = False  # pyright: ignore[reportConstantRedefinition]

disable_numba_log

disable_numba_log()

Disable all formatted Numba JIT/cache status output.

Source code in src/natal/numba/utils.py

def disable_numba_log():
    """Disable all formatted Numba JIT/cache status output."""
    global NUMBA_LOG_ENABLED
    NUMBA_LOG_ENABLED = False  # pyright: ignore[reportConstantRedefinition]

    try:
        from numba.core import config as numba_config  # pyright: ignore
        config_obj: Any = numba_config
        setattr(config_obj, "DEBUG_CACHE", 0)  # noqa: B010
    except Exception:
        pass

disable_numba_signature_trace

disable_numba_signature_trace()

Disable signature diagnostics for newly created Numba specializations.

Source code in src/natal/numba/utils.py

def disable_numba_signature_trace():
    """Disable signature diagnostics for newly created Numba specializations."""
    global NUMBA_SIGNATURE_TRACE_ENABLED
    NUMBA_SIGNATURE_TRACE_ENABLED = False  # pyright: ignore[reportConstantRedefinition]

enable_numba

enable_numba()

Re-enable Numba JIT compilation globally (default state).

Numba is enabled by default for maximum performance. Call this to restore Numba after it was disabled.

Source code in src/natal/numba/utils.py

def enable_numba():
    """Re-enable Numba JIT compilation globally (default state).

    Numba is enabled by default for maximum performance.
    Call this to restore Numba after it was disabled.
    """
    global NUMBA_ENABLED
    NUMBA_ENABLED = True  # pyright: ignore[reportConstantRedefinition]
    _apply_numba_cache_dir()

enable_numba_log

enable_numba_log()

Enable formatted Numba JIT/cache status output (default state).

Source code in src/natal/numba/utils.py

def enable_numba_log():
    """Enable formatted Numba JIT/cache status output (default state)."""
    global NUMBA_LOG_ENABLED
    NUMBA_LOG_ENABLED = True  # pyright: ignore[reportConstantRedefinition]

    try:
        from numba.core import config as numba_config  # pyright: ignore
        config_obj: Any = numba_config
        setattr(config_obj, "DEBUG_CACHE", 1)  # noqa: B010
        _install_cache_log_formatter()
        _install_dispatcher_compile_formatter()
    except Exception:
        pass

enable_numba_signature_trace

enable_numba_signature_trace()

Enable signature diagnostics for newly created Numba specializations.

Source code in src/natal/numba/utils.py

def enable_numba_signature_trace():
    """Enable signature diagnostics for newly created Numba specializations."""
    global NUMBA_SIGNATURE_TRACE_ENABLED
    NUMBA_SIGNATURE_TRACE_ENABLED = True  # pyright: ignore[reportConstantRedefinition]

get_numba_cache_dir

get_numba_cache_dir() -> str

Return the active Numba cache directory.

Source code in src/natal/numba/utils.py

def get_numba_cache_dir() -> str:
    """Return the active Numba cache directory."""
    return NUMBA_CACHE_DIR

is_numba_enabled

is_numba_enabled() -> bool

Check if Numba JIT compilation is currently enabled.

Returns:

Name	Type	Description
bool	bool	True if Numba is enabled (default), False if disabled

Source code in src/natal/numba/utils.py

def is_numba_enabled() -> bool:
    """
    Check if Numba JIT compilation is currently enabled.

    Returns:
        bool: True if Numba is enabled (default), False if disabled
    """
    return NUMBA_ENABLED

is_numba_log_enabled

is_numba_log_enabled() -> bool

Check if Numba JIT/cache status logging is enabled.

Source code in src/natal/numba/utils.py

def is_numba_log_enabled() -> bool:
    """Check if Numba JIT/cache status logging is enabled."""
    return NUMBA_LOG_ENABLED

is_numba_signature_trace_enabled

is_numba_signature_trace_enabled() -> bool

Check if specialization-signature trace logging is enabled.

Source code in src/natal/numba/utils.py

def is_numba_signature_trace_enabled() -> bool:
    """Check if specialization-signature trace logging is enabled."""
    return NUMBA_SIGNATURE_TRACE_ENABLED

njit_switch

njit_switch(func: Callable[P, R], *, cache: bool = True, parallel: bool = False, fastmath: bool = False, **njit_kwargs: Any) -> Callable[P, R]

njit_switch(func: None = None, *, cache: bool = True, parallel: bool = False, fastmath: bool = False, **njit_kwargs: Any) -> Callable[[Callable[P, R]], Callable[P, R]]

njit_switch(func: Optional[Callable[P, R]] = None, *, cache: bool = True, parallel: bool = False, fastmath: bool = False, **njit_kwargs: Any) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]

Numba @njit decorator for functions (controlled by global NUMBA_ENABLED flag).

Parameters:

Name	Type	Description	Default
func	Optional[Callable[P, R]]	Function to decorate	None
cache	bool	Cache compiled functions (default: True)	True
parallel	bool	Enable automatic parallelization (default: False)	False
fastmath	bool	Enable fast math optimizations (default: False)	False
**njit_kwargs	Any	Additional arguments for numba.njit	{}

Examples:

@njit_switch
def my_func(x):
    ...

@njit_switch(parallel=True, fastmath=True)
def my_parallel_func(x):
    ...

Source code in src/natal/numba/utils.py

def njit_switch(
    func: Optional[Callable[P, R]] = None,
    *,
    cache: bool = True,
    parallel: bool = False,
    fastmath: bool = False,
    **njit_kwargs: Any,
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Numba @njit decorator for functions (controlled by global NUMBA_ENABLED flag).

    Args:
        func: Function to decorate
        cache: Cache compiled functions (default: True)
        parallel: Enable automatic parallelization (default: False)
        fastmath: Enable fast math optimizations (default: False)
        **njit_kwargs: Additional arguments for numba.njit

    Examples:
        ```python
        @njit_switch
        def my_func(x):
            ...

        @njit_switch(parallel=True, fastmath=True)
        def my_parallel_func(x):
            ...
        ```
    """

    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
        numba_func: Optional[Callable[P, R]] = None

        if NUMBA_ENABLED:
            try:
                from numba import njit  # pyright: ignore
                from numba.core import config as numba_config  # pyright: ignore

                _apply_numba_cache_dir()

                # Keep all JIT/cache output under one project switch.
                config_obj: Any = numba_config
                setattr(config_obj, "DEBUG_CACHE", 1 if NUMBA_LOG_ENABLED else 0)  # noqa: B010
                _install_cache_log_formatter()
                _install_dispatcher_compile_formatter()

                numba_func = cast(
                    Callable[P, R],
                    njit(
                        fn,
                        cache=cache,
                        parallel=parallel,
                        fastmath=fastmath,
                        **njit_kwargs,
                    ),
                )
            except ImportError:
                pass
            except Exception:
                pass

        # Return compiled or original function
        return numba_func if numba_func is not None else fn

    # Support both @njit_switch and @njit_switch(...)
    if func is not None:
        return decorator(func)
    return decorator

numba_disabled

numba_disabled()

Context manager to temporarily disable Numba (use pure Python).

Useful for debugging within a specific block. Numba is enabled by default, so it will be automatically restored after exiting the context.

Examples:

from natal.numba_utils import numba_disabled

# Default: Numba enabled
with numba_disabled():
    pop.run(n_steps=10)  # Temporarily uses pure Python
# Numba automatically re-enabled here

Source code in src/natal/numba/utils.py

@contextmanager
def numba_disabled():
    """Context manager to temporarily disable Numba (use pure Python).

    Useful for debugging within a specific block. Numba is enabled by default,
    so it will be automatically restored after exiting the context.

    Examples:
        ```python
        from natal.numba_utils import numba_disabled

        # Default: Numba enabled
        with numba_disabled():
            pop.run(n_steps=10)  # Temporarily uses pure Python
        # Numba automatically re-enabled here
        ```
    """
    global NUMBA_ENABLED
    original_state = NUMBA_ENABLED
    NUMBA_ENABLED = False  # pyright: ignore[reportConstantRedefinition]
    try:
        yield
    finally:
        NUMBA_ENABLED = original_state  # pyright: ignore[reportConstantRedefinition]

numba_enabled

numba_enabled()

Context manager to ensure Numba is enabled (default state).

Numba is enabled by default, but this can force re-enable it if temporarily disabled elsewhere. The original state is restored after exiting the context.

Examples:

from natal.numba_utils import numba_enabled

with numba_enabled():
    pop.run(n_steps=100)  # Guaranteed to run with JIT
# Original state is restored after

Source code in src/natal/numba/utils.py

@contextmanager
def numba_enabled():
    """Context manager to ensure Numba is enabled (default state).

    Numba is enabled by default, but this can force re-enable it if temporarily
    disabled elsewhere. The original state is restored after exiting the context.

    Examples:
        ```python
        from natal.numba_utils import numba_enabled

        with numba_enabled():
            pop.run(n_steps=100)  # Guaranteed to run with JIT
        # Original state is restored after
        ```
    """
    global NUMBA_ENABLED
    original_state = NUMBA_ENABLED
    NUMBA_ENABLED = True  # pyright: ignore[reportConstantRedefinition]
    try:
        yield
    finally:
        NUMBA_ENABLED = original_state  # pyright: ignore[reportConstantRedefinition]

with_numba_disabled

with_numba_disabled(func: Callable[P, R]) -> Callable[P, R]

Decorator to run a function with Numba disabled (pure Python).

Useful for testing functions with pure Python, which provides better error messages and debugging capabilities. Numba is enabled by default, so this decorator will restore it after the function returns.

Examples:

from natal.numba_utils import with_numba_disabled

@with_numba_disabled
def debug_version():
    pop.run(n_steps=10)
    # Runs with pure Python for easier debugging

debug_version()  # Temporarily disables Numba
# Numba automatically restored after

Source code in src/natal/numba/utils.py

def with_numba_disabled(func: Callable[P, R]) -> Callable[P, R]:
    """Decorator to run a function with Numba disabled (pure Python).

    Useful for testing functions with pure Python, which provides better
    error messages and debugging capabilities. Numba is enabled by default,
    so this decorator will restore it after the function returns.

    Examples:
        ```python
        from natal.numba_utils import with_numba_disabled

        @with_numba_disabled
        def debug_version():
            pop.run(n_steps=10)
            # Runs with pure Python for easier debugging

        debug_version()  # Temporarily disables Numba
        # Numba automatically restored after
        ```
    """
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        with numba_disabled():
            return func(*args, **kwargs)
    return wrapper

with_numba_enabled

with_numba_enabled(func: Callable[P, R]) -> Callable[P, R]

Decorator to ensure a function runs with Numba enabled (default state).

Since Numba is enabled by default, this is primarily useful for ensuring a function uses JIT even if Numba was temporarily disabled elsewhere.

Examples:

from natal.numba_utils import with_numba_enabled

@with_numba_enabled
def performance_critical():
    pop.run(n_steps=1000)
    # Guaranteed to run with JIT for maximum speed

performance_critical()  # Numba guaranteed enabled
# Original state is restored after

Source code in src/natal/numba/utils.py

def with_numba_enabled(func: Callable[P, R]) -> Callable[P, R]:
    """Decorator to ensure a function runs with Numba enabled (default state).

    Since Numba is enabled by default, this is primarily useful for ensuring
    a function uses JIT even if Numba was temporarily disabled elsewhere.

    Examples:
        ```python
        from natal.numba_utils import with_numba_enabled

        @with_numba_enabled
        def performance_critical():
            pop.run(n_steps=1000)
            # Guaranteed to run with JIT for maximum speed

        performance_critical()  # Numba guaranteed enabled
        # Original state is restored after
        ```
    """
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        with numba_enabled():
            return func(*args, **kwargs)
    return wrapper