Integral Transform

IntegralTransform¤

`neojax.nn.integral_transform.IntegralTransform` ¤

Kernel integral transform (Graph Neural Operator layer).

Approximates the integral of a learned kernel \(k\) over the neighborhood \(A(x)\) of each output point \(x\) by a Monte-Carlo sum (or mean) over the neighboring input points \(y\). Depending on transform_type it computes one of

"linear_kernelonly": \(\int_{A(x)} k(x, y) \, dy\),
"linear": \(\int_{A(x)} k(x, y) \, f(y) \, dy\),
"nonlinear_kernelonly": \(\int_{A(x)} k(x, y, f(y)) \, dy\),
"nonlinear": \(\int_{A(x)} k(x, y, f(y)) \, f(y) \, dy\).

The kernel \(k\) is parametrized by a PointwiseMLP that acts on the concatenated features \((y, x [, f(y)])\) of each pair of points and that only mixes channels. The product with \(f(y)\) is elementwise.

The neighborhoods are passed as padded Neighbors. Optional per-point quadrature weights \(\delta(y)\) turn the sum into a Riemann sum.

With a weighting_fn \(w\) the hard ball \(A(x) = B_r(x)\) is replaced by a smooth weight, i.e. the layer computes \(\int w(x, y) \, k(x, y) \, f(y) \, dy\), the mollified GNO. This makes the output differentiable in the query position \(x\).

Parameters:

Name	Type	Description	Default
`channel_mlp`	`PointwiseMLP`	The `PointwiseMLP` parametrizing the kernel. Its input channels must be the sum of the channels of \(y\) and \(x\) (and of \(f\) for the nonlinear transforms). For the `"linear"` and `"nonlinear"` transforms its output channels must equal the channels of \(f\).	required
`transform_type`	`Literal['linear_kernelonly', 'linear', 'nonlinear_kernelonly', 'nonlinear']`	Integral transform to compute. Default is `"linear"`.	`'linear'`
`reduction`	`Literal['sum', 'mean']`	How to aggregate over a neighborhood, by summing (`"sum"`) or averaging (`"mean"`). Default is `"sum"`. An empty neighborhood yields zero for both. Forced to `"sum"` if `weighting_fn` is given.	`'sum'`
`weighting_fn`	`Callable[[Float[Array, ...]], Float[Array, ...]] \| None`	Optional function mapping the squared distances between \(x\) and its neighbors to weights by which the kernel values are multiplied. This is the mollifier of the mollified GNO, see `make_weighting_fn`. The squared distances are recomputed by the `GNOBlock` from the coordinates, so the weights stay differentiable in \(x\) even for precomputed neighbors. Default is `None`.	`None`
`query_chunk_size`	`int \| None`	If given, the transform is computed for this many output points at a time and rematerialized in the backward pass. The kernel `channel_mlp` is evaluated on (num pairs) = `m * k` pairs, which can be prohibitively large in memory for wide kernels, so this trades speed for memory. Default is `None`.	`None`

Raises:

Type	Description
`ValueError`	If `transform_type` or `reduction` is invalid.
`ValueError`	If `query_chunk_size` is not a positive integer.

Internal Attributes

channel_mlp (PointwiseMLP): The kernel network.
transform_type (str): Integral transform type.
reduction (str): Reduction over neighborhoods.
weighting_fn (Callable | None): Weighting function of the squared distances.
query_chunk_size (int | None): Number of output points per chunk.

Sum or mean with a weighting_fn

Weights near the boundary of a neighborhood are close to zero, but such neighbors still count in a "mean". When the query moves, neighbors enter and leave the ball and the mean changes discontinuously. That is why a weighting_fn forces the "sum" reduction.

Neighbor truncation

Neighbors beyond the max_neighbors nearest are dropped, even if they have a nonzero weight. This reintroduces a discontinuity, so choose max_neighbors large enough that truncation never occurs when using a weighting_fn for derivatives.

`call(y: Float[Array, 'n e_y'], x: Float[Array, 'm e_x'], neighbors: Neighbors, f_y: Float[Array, 'f_c n'] | None = None, quadrature_weights: Float[Array, n] | None = None) -> Float[Array, 'out_c m']` ¤

Computes the kernel integral transform.

Parameters:

Name	Type	Description	Default
`y`	`Float[Array, 'n e_y']`	The (embedded) coordinates of the `n` input points that are integrated over.	required
`x`	`Float[Array, 'm e_x']`	The (embedded) coordinates of the `m` output points.	required
`neighbors`	`Neighbors`	The neighbors in `y` of each output point, shaped (m, k). These must be found with the unembedded coordinates, e.g. by `radius_neighbor_search`.	required
`f_y`	`Float[Array, 'f_c n'] \| None`	Input function sampled at `y`. Optional for the `"linear"` transform, ignored by `"linear_kernelonly"` and required by the nonlinear transforms.	`None`
`quadrature_weights`	`Float[Array, n] \| None`	Optional quadrature weights \(\delta(y)\) of the `n` input points, that multiply the kernel values (Riemann sum). With a `"sum"` reduction they make the output independent of the discretization. Default is `None`, no quadrature weights.	`None`

Returns:

Type	Description
`Float[Array, 'out_c m']`	The output function sampled at `x`.

Raises:

Type	Description
`ValueError`	If `f_y` is required but `None`.
`ValueError`	If the `"linear"` or `"nonlinear"` transforms are used and the kernel output channels do not match the channels of `f_y`.

Integral Transform