Mollifier Functions

Mollifier Functions¤

Mollifier functions, also known as weighting functions, are optionally used in the output Graph Neural Operator (GNO) of the GiNO. They replace the hard indicator function over the ball \(\mathbb{1}_{B_r(x)}\) in the kernel integral of the GNO with a smooth differentiable weighting function.

Cite

Enabling Automatic Differentiation with Mollified Graph Neural Operators

@article{lin2025enabling,
    title={Enabling automatic differentiation with mollified graph neural operators},
    author={Lin, Ryan Y and Berner, Julius and Duruisseaux, Valentin and Pitt, David and Leibovici, Daniel and Kossaifi, Jean and Azizzadenesheli, Kamyar and Anandkumar, Anima},
    journal={arXiv preprint arXiv:2504.08277},
    year={2025}
}

Squared distances

The weights are evaluated on squared distances and take the squared radius, i.e. \(d = \Vert x - y \Vert^2 / r^2\). This follows neuraloperator and the pseudocode of the mGNO paper. The paper's main text writes the same functions with \(d = \Vert x - y \Vert / r\). The quadr_cutoff weighting function is not in the paper.

Neighbor truncation and Mollifiers

Neighbors beyond max_neighbors are dropped in the neighbor search, even with a nonzero weight, which reintroduces a discontinuity. Choose max_neighbors large enough that truncation never occurs if you rely on derivatives.

`neojax.nn.mollifier_functions.bump_cutoff(x: Float[Array, ...], radius: float = 1.0, scale: float = 1.0, eps: float = 1e-07) -> Float[Array, ...]` ¤

Bump cutoff function with exponential decay.

\(w(x) = s \cdot e \cdot \exp\left(-1 / (1 - d^2 + \varepsilon)\right)\) with \(d = x / r\) clipped to \([0, 1]\) where \(e\) is Euler's number.

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, ...]`	Values in \([0, r]\).	required
`radius`	`float`	Cutoff radius \(r\).	`1.0`
`scale`	`float`	Scaling factor \(s\) of the weights.	`1.0`
`eps`	`float`	Small constant for numerical stability.	`1e-07`

Returns:

Type	Description
`Float[Array, ...]`	The weights, same shape as `x`.

`neojax.nn.mollifier_functions.half_cos_cutoff(x: Float[Array, ...], radius: float = 1.0, scale: float = 1.0) -> Float[Array, ...]` ¤

Half-cosine cutoff function with smooth decay.

\(w(x) = s \, (0.5 \cos(\pi d) + 0.5)\) with \(d = x / r\).

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, ...]`	Values in \([0, r]\).	required
`radius`	`float`	Cutoff radius \(r\).	`1.0`
`scale`	`float`	Scaling factor \(s\) of the weights.	`1.0`

Returns:

Type	Description
`Float[Array, ...]`	The weights, same shape as `x`.

`neojax.nn.mollifier_functions.quadr_cutoff(x: Float[Array, ...], radius: float = 1.0, scale: float = 1.0) -> Float[Array, ...]` ¤

Piecewise quadratic cutoff function.

\(w(x) = s \cdot \begin{cases} 1 - 2d^2 & d < 0.5 \\ 2(1 - d)^2 & d \geq 0.5 \end{cases}\) with \(d = x / r\).

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, ...]`	Values in \([0, r]\).	required
`radius`	`float`	Cutoff radius \(r\).	`1.0`
`scale`	`float`	Scaling factor \(s\) of the weights.	`1.0`

Returns:

Type	Description
`Float[Array, ...]`	The weights, same shape as `x`.

`neojax.nn.mollifier_functions.quartic_cutoff(x: Float[Array, ...], radius: float = 1.0, scale: float = 1.0) -> Float[Array, ...]` ¤

Quartic cutoff function.

\(w(x) = s \, (1 - d^2)^2\) with \(d = x / r\).

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, ...]`	Values in \([0, r]\).	required
`radius`	`float`	Cutoff radius \(r\).	`1.0`
`scale`	`float`	Scaling factor \(s\) of the weights.	`1.0`

Returns:

Type	Description
`Float[Array, ...]`	The weights, same shape as `x`.

`neojax.nn.mollifier_functions.octic_cutoff(x: Float[Array, ...], radius: float = 1.0, scale: float = 1.0) -> Float[Array, ...]` ¤

Octic cutoff function with eighth-order polynomial.

\(w(x) = s \, (-3d^8 + 8d^6 - 6d^4 + 1)\) with \(d = x / r\).

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, ...]`	Values in \([0, r]\).	required
`radius`	`float`	Cutoff radius \(r\).	`1.0`
`scale`	`float`	Scaling factor \(s\) of the weights.	`1.0`

Returns:

Type	Description
`Float[Array, ...]`	The weights, same shape as `x`.

`neojax.nn.mollifier_functions.make_weighting_fn(name: typing.Literal['bump', 'half_cos', 'quadr', 'quartic', 'octic'], sq_radius: float, scale: float = 1.0) -> collections.abc.Callable[[Float[Array, ...]], Float[Array, ...]]` ¤

Selects a weighting function for the output GNO of a mollified GNO.

The weighting functions take the squared radius of the neighborhoods because they are evaluated on the squared distances between query and neighbor. They replace the hard cutoff of the neighborhood with a smooth weight, which makes the GNO differentiable in the query position.

Parameters:

Name	Type	Description	Default
`name`	`Literal['bump', 'half_cos', 'quadr', 'quartic', 'octic']`	Name of the weighting function. One of `"bump"`, `"half_cos"`, `"quadr"`, `"quartic"` or `"octic"`.	required
`sq_radius`	`float`	Squared radius of the neighborhoods.	required
`scale`	`float`	Factor by which to scale all weights. Default is 1.0.	`1.0`

Returns:

Type	Description
`Callable[[Float[Array, ...]], Float[Array, ...]]`	A function mapping squared distances to weights.

Raises:

Type	Description
`ValueError`	If `name` is not a registered weighting function.

Mollifier Functions