Source code for vivyd.core.taichi

from ._taichi_fallback import TaichiFallback

from typing import Protocol, runtime_checkable

try:
    import taichi as ti
except ImportError:
    _IS_TAICHI_AVAILABLE = False
    ti_or_fallback = TaichiFallback()
else:
    _IS_TAICHI_AVAILABLE = True
    ti_or_fallback = ti

ti_or_fallback.__doc__ = ("""
    If Taichi is available, this is the Taichi module. Otherwise, this is a
    fallback object that raises errors when Taichi features are accessed.
""")


_USE_TAICHI = False




def use_taichi(use_taichi: bool=True, **ti_params):
    """
    Enable or disable Taichi acceleration for compatible objects.
    
    Parameters
    ----------
    use_taichi : bool, optional
        Whether to enable Taichi acceleration. Default is `True`.
    **ti_params
        Additional parameters to pass to `ti.init()` when enabling Taichi.
        Most commonly used parameters include:

        - `arch`: The Taichi backend to use (e.g., `ti.cpu`, `ti.gpu`, etc.). Default is `ti.cpu`.
        - `cpu_max_num_threads`: The maximum number of CPU threads to use when `arch` is `ti.cpu`.
        
        The default parameters are also modified to use double precision floating point (`ti.f64`)
        and to disable fast math optimizations, which can improve numerical accuracy at the cost of performance.

        See `<https://docs.taichi-lang.org/docs/hello_world#initialize-taichi>`_ for more information.
    """
    
    global _USE_TAICHI

    if use_taichi:
        
        if _IS_TAICHI_AVAILABLE:
            _USE_TAICHI = True
            
            if "arch" not in ti_params:
                ti_params["arch"] = ti_or_fallback.cpu
            if "cpu_max_num_threads" not in ti_params:
                ti_params["cpu_max_num_threads"] = 1
            if "default_fp" not in ti_params:
                ti_params["default_fp"] = ti_or_fallback.f64
            if "fast_math" not in ti_params:
                ti_params["fast_math"] = False
            
            print("Taichi is available. Initializing Taichi with parameters:", ti_params)

            ti_or_fallback.init(**ti_params)
        
        if not _IS_TAICHI_AVAILABLE:
            _USE_TAICHI = False
            print ("Taichi is not available. Falling back to classic Python implementation.")
            
    else:
        _USE_TAICHI = False

        
        


def is_taichi_used() -> bool:
    """
    Check whether Taichi acceleration is currently enabled.
    
    Returns
    -------
    bool
        `True` if Taichi acceleration is enabled, `False` otherwise.
    """    
    
    return _USE_TAICHI





class TaichiContext:
    """
    A context manager for temporarily enabling Taichi acceleration within
    a block of code.
    
    Parameters
    ----------
    **ti_params
        Additional parameters to pass to `ti.init()` when enabling Taichi.
        See :func:`use_taichi` for more details on the available parameters.
    
    Example
    -------
    .. code-block:: python
    
        with TaichiContext():
            # Code in this block will use Taichi acceleration if available
            ...
    """
    
    def __init__(self, **ti_params):
        self.ti_params = ti_params
    
    def __enter__(self):
        use_taichi(True, **self.ti_params)
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        use_taichi(False)

        
    


@runtime_checkable
class TaichiCompatible(Protocol):
    """
    A protocol for objects that are compatible with Taichi acceleration.
    """
    


    def __call__(self, *args, **kwargs): ...


    def _call_taichi_func(self, *args, **kwargs): ...
    
    def _call_taichi_kernel(self, *args, **kwargs): ...
    
    def _call_taichi(self, *args, **kwargs): ...

    



@runtime_checkable
class TaichiConvertible(Protocol):
    """A protocol for objects that can be converted to a Taichi-compatible form."""



    def to_taichi(self, *args, **kwargs) -> TaichiCompatible:
        """Convert the object to a Taichi-compatible object."""
        ...