brepgraph

Turn B-rep CAD files into ML-ready graphs. A STEP file goes in; a framework-neutral NumPy graph comes out. Each face becomes a node (surface type, area, centroid, a UV grid of points, normals and trim mask) and each shared edge becomes a graph edge (curve type, length, convexity, dihedral angle, a sampled curve grid). Thin adapters hand the graph to PyTorch Geometric or DGL.

• pip install only: built on the OCP wheels from PyPI. No conda, no proprietary CAD kernel.
• Featurization only: no models, no training code.
• Features are registered functions, so you can add or mix your own (docs/extending.md).
• Interactive Plotly views of the part, its face graph, and individual faces for debugging features.

Installation

pip install brepgraph               # numpy + OCP
pip install "brepgraph[pyg]"        # + PyTorch Geometric adapter
pip install "brepgraph[dgl]"        # + DGL adapter
pip install "brepgraph[viz]"        # + Plotly visualization

Python 3.10 to 3.13 on Linux, macOS and Windows. PyTorch and DGL builds are CUDA- and version-specific, so they are not pinned; install the build that matches your machine first.

Quick start

import brepgraph as bg

g = bg.load("part.step")                      # default `uvnet` feature set
print(g)                                      # every feature with its shape and dtype
g.num_faces, g.num_edges
g.node_feats["uv_grid"].shape                 # (F, 7, 10, 10): x y z nx ny nz mask
g.edge_feats["u_grid"].shape                  # (E, 6, 10):     x y z tx ty tz

g.save("part.npz")                            # plain NumPy archive, no pickle
g = bg.BrepGraph.load("part.npz")

data = g.to_pyg(x=("area", "centroid"))       # torch_geometric.data.Data, both edge directions
dgl_graph = g.to_dgl()                        # DGL graph

g = bg.load("part.step", "brepmfr")           # BrepMFR-style set incl. face-pair D2 / A3 histograms
g.pair_feats["d2"].shape                      # (F, F, 64)

brepgraph convert part.step part.npz --preset uvnet
brepgraph info part.npz
brepgraph view part.step --color-by area      # interactive HTML next to the input
brepgraph view part.step --kind graph         # the face graph instead of the 3-D part
brepgraph convert assembly.step asm.npz --all-solids   # one graph per solid
brepgraph batch data/ -o graphs/              # a whole dataset in parallel, with report.json

All errors derive from bg.BrepGraphError (FileReadError, UnsupportedShapeError, InvalidSolidError, FeatureError, MissingExtraError), so a batch job only needs one try/except:

for path in files:
    try:
        bg.load(path).save(path.with_suffix(".npz"))
    except bg.BrepGraphError as exc:
        print(path.name, "skipped:", exc)

Datasets and multi-body files

graphs = bg.load_all("assembly.step")          # one graph per solid; meta["solid_index"] says which

results = bg.load_many(paths, workers=8)       # parallel worker processes, results in input order
ok = [r.graphs[0] for r in results if r.ok]    # failed files carry r.error instead of raising

brepgraph batch data/ -o graphs/ --workers 8 --all-solids   # mirrors folders, writes graphs/report.json

Worker processes are used because OpenCASCADE is not thread-safe. The default is min(CPUs, 8) workers; each one holds a full OpenCASCADE runtime (a few hundred MB), so lower workers on machines with little free memory. If OpenCASCADE crashes on a corrupt file, only that file is reported as failed; the rest of the batch continues. Custom features must live in an importable module to be visible to workers (in a notebook, use workers=1).

Features

Node, edge and pair features are stored as separate named arrays (g.node_feats, g.edge_feats, g.pair_feats). Row i of a node array belongs to face i; row k of an edge array belongs to column k of g.edge_index. Categorical features are integer indices.

The brepmfr preset follows Zhang et al., BrepMFR, CAGD 111 (2024): 5x5 grids plus face features num_loops and num_adjacent_faces, and face-pair matrices shortest_path [F, F], d2 and a3 [F, F, 64] (frequencies; rows sum to 1). The paper does not say how the three A3 points are drawn from two faces; the choice made here is documented in src/brepgraph/features/pair.py. Pair features cost O(F^2) in time and memory, so they are opt-in: as float32 one [F, F, 64] histogram takes about 200 MB at 900 faces, and d2 + a3 take about two minutes for a 900-face part. ("d2", {"dtype": "uint16"}) stores the paper's exact counts and "float16" stores frequencies, both at half the memory.

• Complete reference (generated from the registry): docs/features.md
• Graph layout and conventions (normals, mask, seam edges, ordering, units): docs/schema.md
• Writing your own feature: docs/extending.md
• Feeding the graph to a GNN: docs/using-with-gnns.md and examples/gnn_with_pyg.py
• Demo notebook: notebooks/demo.ipynb

Visualization

bg.viz.show(g, "surface_type").show()         # part coloured by a feature, edges coloured by convexity
bg.viz.show(g, "area", graph=True).show()     # continuous colours, plus the adjacency graph overlay
bg.viz.show_graph(g, layout="centroid").show()  # the face graph on its own
bg.viz.show_face(g, 12).show()                # one face's UV samples, trim mask and normals

show re-loads the source file recorded in the graph and checks it against the stored areas and centroids, so a changed face order is reported instead of silently mis-coloured. Archives do not store the absolute input path by default (it would leak local directory names); for a graph loaded from an archive pass source=path, or save with include_source_path=True.

Known limitations

• bg.load accepts exactly one solid; use bg.load_all (CLI --all-solids) for assemblies and multi-body files. Open shells and non-manifold edges raise UnsupportedShapeError.
• Seam and degenerated edges are dropped (a sphere is 1 face and 0 edges). This is not configurable yet.
• Units are whatever OCCT produces (millimetres for STEP); nothing is rescaled unless normalize=True.
• Speed: every sample is a Python-to-C++ call, roughly 20-25 ms per face per process with the default 10x10 grid. load_many / brepgraph batch use all CPU cores.
• OpenCASCADE has global state: parallelise with processes (load_many), not threads.
• Face and edge order is deterministic for a given file and OCP version, but may differ between OCP versions.
• A face graph can be disconnected (for example a solid with an internal cavity); unreachable pairs get shortest_path = -1.
• Not included: coedge graphs, an on-disk cache, per-file timeouts (a file that hangs blocks its worker).
• The DGL adapter was tested against a stub of the DGL API because no native DGL build was available on the development machine; run pytest tests/test_v02.py where dgl imports to exercise the real thing.

Development

pip install -e ".[dev]"
ruff check src tests scripts examples && ruff format --check src tests scripts examples
mypy
pytest                       # fast suite;  `pytest -m slow` adds a 535-face part

Features are tested against shapes with known answers (box, cylinder, sphere, cone, torus, pocketed and filleted boxes, a plate with a hole), against invariants on every input (unit normals where the mask is 1, areas that sum to the solid's surface area, rigid-motion and scale behaviour, determinism, round trips), and on harder generated parts. Supported OCP lines are 7.9 (Python 3.10) and 8.x (Python 3.11+). See CONTRIBUTING.md.

Citing

If this package is useful in your research, please cite it (see CITATION.cff) together with the papers whose feature definitions it follows.

Credits and license

The feature ideas follow UV-Net (Jayaraman et al., CVPR 2021), BRepNet (Lambourne et al., CVPR 2021) and BrepMFR (Zhang et al., CAGD 2024). This package is written independently against the papers and reuses none of their code. Geometry is by OpenCASCADE Technology (LGPL 2.1 with exception) through the OCP bindings (cadquery-ocp), installed as a separate dependency.

Released under the MIT License.