Source code for nvalchemiops.math.math
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any
import warp as wp
@wp.func
def wp_safe_divide(x: wp.float64, y: wp.float64) -> wp.float64:
"""Safe division.
Divides x by y, with a safe division to avoid division by zero.
"""
return wp.where(y < wp.float64(1e-8), wp.float64(0.0), x / y)
@wp.func
def wp_exp_kernel(x: wp.float64, factor: wp.float64) -> wp.float64:
"""
Safe exponential multiplication and division.
Calculates exp(-x * factor) / x, with a safe division to avoid division by zero.
"""
return wp_safe_divide(wp.exp(-x * factor), x)
@wp.func
def wp_logistic(x: Any) -> Any:
"""Logistic function, evaluated so the exponent is never positive.
The naive ``1 / (1 + exp(-x))`` overflows for large negative ``x``. Harmless in float64,
a real failure in float32, so each branch exponentiates ``-|x|`` instead.
Parameters
----------
x : Any
Input value
Returns
-------
Any
``1 / (1 + exp(-x))``
"""
one = type(x)(1.0)
decay = wp.exp(-wp.abs(x))
if x >= type(x)(0.0):
return one / (one + decay)
return decay / (one + decay)
[docs]
@wp.func
def wpdivmod(a: int, b: int): # type: ignore
"""Integer quotient and remainder.
Computes ``quotient = int(a / b)`` with truncating division toward zero,
then adjusts a negative C-style remainder. For the positive-divisor
contract used by the neighbor kernels, this matches :func:`divmod`,
including for negative dividends (for example, ``wpdivmod(-5, 3)`` and
Python ``divmod(-5, 3)`` both return ``(-2, 1)``).
Parameters
----------
a : int
Dividend.
b : int
Positive divisor.
Returns
-------
quotient : int
Floor quotient for the positive divisor.
remainder : int
Non-negative remainder in ``[0, |b|)`` produced by the adjustment
above.
"""
div = int(a / b)
mod = a % b
if mod < 0:
div -= 1
mod = b + mod
return div, mod
[docs]
@wp.func
def wp_erfc_f32_fast(x: Any) -> Any:
"""Complementary error function approximation for float32.
Uses the Abramowitz and Stegun approximation with maximum error ~1.5e-7.
erfc(x) = 1 - erf(x) for x >= 0, and erfc(-x) = 2 - erfc(x) for x < 0.
This is a float32-grade approximation: its own error (~1.5e-7) is close to
float32 epsilon, so evaluating it in float64 buys no accuracy, only cost.
Do not use it for float64 inputs -- use native ``wp.erfc`` there, or call
:func:`wp_erfc_input_precision` to dispatch automatically.
Parameters
----------
x : Any
Input value
Returns
-------
Any
erfc(x) approximation
"""
abs_x = wp.abs(x)
# Abramowitz and Stegun constants for erfc approximation
p = type(x)(0.3275911)
a1 = type(x)(0.254829592)
a2 = type(x)(-0.284496736)
a3 = type(x)(1.421413741)
a4 = type(x)(-1.453152027)
a5 = type(x)(1.061405429)
# Compute approximation for |x|
t = type(x)(1.0) / (type(x)(1.0) + p * abs_x)
t2 = t * t
t3 = t2 * t
t4 = t3 * t
t5 = t4 * t
# Polynomial approximation
poly = a1 * t + a2 * t2 + a3 * t3 + a4 * t4 + a5 * t5
# Apply exponential factor
exp_neg_x2 = wp.exp(-abs_x * abs_x)
erfc_abs_x = poly * exp_neg_x2
# Handle sign: erfc(-x) = 2 - erfc(x)
return wp.where(x >= type(x)(0.0), erfc_abs_x, type(x)(2.0) - erfc_abs_x)
[docs]
@wp.func
def wp_erfc(x: Any) -> Any:
"""Deprecated alias for :func:`wp_erfc_f32_fast`.
.. deprecated::
Kept only for backward compatibility: existing callers that evaluate
this float32-grade approximation at float64 get exactly the same
(imprecise, for float64) result they always did. New code should call
:func:`wp_erfc_input_precision`, which dispatches to native
``wp.erfc`` for float64 and to :func:`wp_erfc_f32_fast` otherwise.
"""
return wp_erfc_f32_fast(x)
@wp.func
def wp_erfc_input_precision(x: wp.float64) -> wp.float64:
"""``erfc(x)`` at the precision of ``x``: native for float64, fast for float32.
Native ``wp.erfc`` is exact to machine precision but costs more than the
~1.5e-7-accurate polynomial in :func:`wp_erfc_f32_fast`, which is already
below float32 epsilon. Use this dispatcher in code that evaluates in the
caller's input precision rather than a type hard-coded to one dtype.
"""
return wp.erfc(x)
