Neural-operator data

scimba_jax.neural_operator.data_for_no contains the geometry and topology containers used by neural operators. They build sparse connectivity on the host, then keep numerical arrays as JAX pytree leaves. Consequently, topology is not rebuilt inside a JIT trace and changed geometry with the same array shapes can reuse an executable.

Choose the representation

Data	Use it for	Connectivity	Quadrature information
GridData	Tensor-product, meshless grids	implicit grid stencil	spacings in grid_size
PointCloudData	Samples without a mesh	compact CSR radius graph	density and integration_weights
MeshData(entity="cells")	Cell-centred fields	cells sharing a face	cell measures
MeshData(entity="nodes")	Nodal fields	geometric mesh edges	lumped dual-cell masses

GridData stores the coordinates of a tensor-product grid and its physical step in every direction. PointCloudData stores a radius-neighbour graph in CSR form: neighbor_offsets[i]:neighbor_offsets[i + 1] indexes the neighbours of point i; senders and receivers contain the corresponding directed edges. This is linear in the number of edges, unlike a padded N x N graph.

Mesh graphs

MeshData accepts Mesh, UnstructuredMesh, and BlockStructuredMesh. Cell graphs connect cells through interior faces, including block interfaces. Nodal graphs use Q1 vertices for structured and block meshes, and all geometric nodes for unstructured meshes. Coincident vertices on conforming block interfaces are merged. Non-conforming interfaces deliberately remain separate nodes: their weighted mortar/interpolation links are finite-element transfers, not ordinary graph edges.

from scimba_jax.neural_operator.data_for_no.mesh_data import MeshData

cell_graph = MeshData(dim=2, mesh=mesh, entity="cells")
nodal_graph = MeshData(dim=2, mesh=mesh, entity="nodes")
integral = cell_graph.integrate(cell_values)

For a cell graph, integration_weights() returns physical cell measures. For a nodal graph, it returns a positive lumped mass obtained from element quadrature, i.e. the dual-cell volume. density is a multiplicative physical weight and defaults to one. sampling_density() is instead the probability density induced by the discrete quadrature weights; it is useful when a loss uses importance sampling.

Multilevel pooling

HierarchicMeshData and HierarchicPointCloudData order levels from coarse to fine. Each transition stores both a fine-to-coarse parent_of vector and the reverse relation in CSR form. This makes pooling and unpooling simple gather/scatter operations, with no nearest-neighbour query during training.

from scimba_jax.neural_operator.data_for_no.hierarchic_mesh_data import (
    HierarchicMeshData,
)

hierarchy = HierarchicMeshData.from_block_meshes([coarse_mesh, fine_mesh])
coarse_features = hierarchy.restrict_mean(0, fine_features)
reconstructed = hierarchy.prolong(0, coarse_features)

parent = hierarchy.parent_of(level=1, cell=fine_cell)
children = hierarchy.children_of(level=0, cell=parent)

For BlockStructuredMesh, levels must share the same macro-patch layout; each fine patch has an integral refinement factor in every direction. The maps are then exact in tensor coordinates, including for curved patches. For a Gmsh domain, use HierarchicMeshData.from_mesh_hierarchy(...) with a MeshHierarchy, or its from_gmsh_points / from_gmsh_curve constructors. The coarsest mesh is generated once and finer levels are nested refinements, so the transfers are exact.

HierarchicPointCloudData.from_fine(...) builds the analogous hierarchy by repeated, deterministic widest-axis median splits. It aims for approximately N / 2**d coarse points at each level, keeps every level’s CSR graph, and preserves total quadrature mass by summing child integration weights.

Public operations

Both hierarchy classes provide:

• parent_of(level, index) and children_of(level, index) for one relation;

• transfer_info(transition) for the parent vector and CSR children together;

• restrict_mean(transition, values) for mean pooling;

• prolong(transition, values) to copy coarse features to their children.

These are topological pooling operations. Finite-element restriction and prolongation for nodal degrees of freedom remain the responsibility of NestedTransfer, because they require interpolation weights rather than just parent/child indices.