Skip to content

MurrellGroup/FuxGNN.jl

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

2 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

FuxGNN.jl

FuxGNN.jl is a minimal, dependency-light graph message passing core designed for ProteinMPNN-style workloads.

It intentionally avoids heavyweight graph frameworks and works directly with explicit source/destination edge index arrays.

Install

using Pkg
Pkg.add(url = "https://github.com/MurrellGroup/FuxGNN.jl")

Core API

  • EdgeIndex(src, dst, num_nodes)
  • num_edges(edge_index)
  • gather_nodes(node_features, idx)
  • scatter_sum(values, dst, num_nodes)
  • scatter_mean(values, dst, num_nodes)
  • cat_neighbor_nodes(node_features, edge_features, edge_index)
  • propagate(node_features, edge_features, edge_index; message_fn, reduce=:sum|:mean)

Quick example

using FuxGNN

# 3 nodes, 3 directed edges.
ei = EdgeIndex([1, 3, 2], [2, 2, 1], 3)

node_features = Float32[
    1 0;
    0 1;
    1 1;
]
edge_features = Float32[
    1 1;
    2 2;
    3 3;
]

message_fn = (src, dst, edge, eidx) -> src .+ dst .+ edge

out = propagate(node_features, edge_features, ei; message_fn = message_fn, reduce = :sum)

Writing a custom layer

A minimal message-passing layer can be written as:

using FuxGNN

function simple_mp_layer(node_features::AbstractMatrix, edge_features::AbstractMatrix, ei::EdgeIndex)
    msg_fn = (src, dst, edge, eidx) -> begin
        # Example: linear-ish residual message.
        0.7f0 .* src .+ 0.3f0 .* dst .+ edge
    end

    agg = propagate(node_features, edge_features, ei; message_fn = msg_fn, reduce = :mean)
    return node_features .+ agg
end

For a full walkthrough, see docs/src/custom_layer.md.

Docs

Multi-page docs are in docs/src:

Build docs locally:

julia docs/make.jl

Tests

julia --project -e 'using Pkg; Pkg.test()'

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages