liblaf.melon.utils ¶

def filter_kwargs(
    func: Callable[..., Any], kwargs: Mapping[str, Any]
) -> Mapping[str, Any]:
    """Filter keyword arguments to those accepted by `func`.

    Functions with `**kwargs` receive the original mapping unchanged.

    Args:
        func: Callable whose signature should be respected.
        kwargs: Candidate keyword arguments.

    Returns:
        Supported keyword arguments.
    """
    from inspect import Parameter

    signature: inspect.Signature = inspect.signature(func)
    filtered: dict[str, Any] = {}
    for p in signature.parameters.values():
        match p.kind:
            case Parameter.POSITIONAL_OR_KEYWORD | Parameter.KEYWORD_ONLY:
                if p.name in kwargs:
                    filtered[p.name] = kwargs[p.name]
            case Parameter.VAR_KEYWORD:
                return kwargs
    return filtered

def pick[K, V](allowlist: Container[K], dictionary: Mapping[K, V]) -> dict[K, V]:
    """Return dictionary items whose keys are in `allowlist`.

    Args:
        allowlist: Keys to keep.
        dictionary: Source mapping.

    Returns:
        A plain dictionary containing only allowed keys.
    """
    return tlz.keyfilter(lambda k: k in allowlist, dictionary)

@contextlib.contextmanager
def temporary_array(
    attributes: pv.DataSetAttributes,
    data: Float[ArrayLike, "..."] | None = None,
    name: str = "",
    length: int = 8,
) -> Generator[str]:
    """Attach a temporary PyVista data array and remove it on exit.

    Args:
        attributes: Point, cell, or field data container to mutate.
        data: Optional array value to store under the generated name.
        name: Prefix for the generated array name.
        length: Number of random suffix characters to append.

    Yields:
        The generated array name.
    """
    suffix: str = "".join(random.choices(_ALPHABET, k=length))  # noqa: S311
    name: str = f"{name}{suffix}"
    if data is not None:
        attributes[name] = data
    try:
        yield name
    finally:
        attributes.pop(name, None)

def warp_stream_from_torch() -> wp.Stream | None:
    if not torch.cuda.is_available():
        return None
    torch_stream: torch.cuda.Stream = torch.cuda.current_stream()
    warp_stream: wp.Stream = wp.stream_from_torch(torch_stream)
    return warp_stream