Skip to content

Commit da4a195

Browse files
Uneeb808Uneeb808
authored
Add ZINC molecular graph dataset (#250)
1 parent 259d2eb commit da4a195

4 files changed

Lines changed: 226 additions & 0 deletions

File tree

‎docs/src/datasets/graphs.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,4 +38,5 @@ TUDataset
3838
WindMillEnergy
3939
AmazonComputers
4040
AmazonPhoto
41+
ZINC
4142
```

‎src/MLDatasets.jl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,9 @@ export WindMillEnergy
147147
include("datasets/graphs/amazon.jl")
148148
export AmazonComputers
149149
export AmazonPhoto
150+
151+
include("datasets/graphs/zinc.jl")
152+
export ZINC
150153
# Meshes
151154

152155
include("datasets/meshes/faust.jl")
@@ -172,6 +175,7 @@ function __init__()
172175
__init__temporalbrains()
173176
__init__windmillenergy()
174177
__init__amazon()
178+
__init__zinc()
175179
# misc
176180
__init__iris()
177181
__init__mutagenesis()

‎src/datasets/graphs/zinc.jl‎

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
function __init__zinc()
2+
DEPNAME = "ZINC"
3+
DOCS = "https://pubs.acs.org/doi/abs/10.1021/acs.jcim.5b00559"
4+
SPLIT_LINK = "https://raw.githubusercontent.com/graphdeeplearning/" *
5+
"benchmarking-gnns/master/data/molecules"
6+
# NOTE for maintainer: NPZ files are currently hosted on the contributor's fork.
7+
# Please re-upload to JuliaML/MLDatasets.jl releases and update this link on merge.
8+
NPZ_LINK = "https://github.com/Uneeb808/MLDatasets.jl/" *
9+
"releases/download/zinc-data"
10+
11+
register(DataDep(
12+
DEPNAME,
13+
"""
14+
Dataset: ZINC Molecular Graph Dataset
15+
Website: $DOCS
16+
17+
~250,000 molecular graphs for graph-level regression of penalized logP
18+
(y = logP - SAS - cycles). Includes a 12k benchmark subset used in
19+
"Benchmarking Graph Neural Networks" (Dwivedi et al. 2020).
20+
21+
Please cite:
22+
Irwin et al. (2012) https://pubs.acs.org/doi/abs/10.1021/acs.jcim.5b00559
23+
Gomez-Bombarelli et al. https://arxiv.org/abs/1610.02415
24+
Dwivedi et al. (2020) https://arxiv.org/abs/2003.00982
25+
""",
26+
[
27+
"$NPZ_LINK/train.npz",
28+
"$NPZ_LINK/val.npz",
29+
"$NPZ_LINK/test.npz",
30+
"$SPLIT_LINK/train.index",
31+
"$SPLIT_LINK/val.index",
32+
"$SPLIT_LINK/test.index",
33+
],
34+
[
35+
"139abb0fb3ce4305c2c04e3c6f55e771022bcf392b4dbb7fb315690f56cd2a96",
36+
"e2d5dad38bd2bff0e2559b463420fd570d290c80910f58b90bd3c7bead8c1149",
37+
"27b61afc6a660871cd7054bc0a29e6240d9c25290e3ebf1a4bf0d320aa906327",
38+
"575d6acd72ac207a95947e2bdd16411ef6fc3faa88c2f55a2c496ea43022a6d8",
39+
"fe48c38157f0d38e7fc441903a143392eda080511a79e8fb8f60f5c1b0a9c164",
40+
"6aa50d98976044fb36089afef5b4c991b429ca1c8e009d33a0319e42a2d9b525",
41+
]
42+
))
43+
end
44+
45+
46+
"""
47+
ZINC(; split=:train, subset=false, dir=nothing)
48+
49+
The ZINC dataset from the [ZINC database](https://pubs.acs.org/doi/abs/10.1021/acs.jcim.5b00559)
50+
and the [Automatic Chemical Design](https://arxiv.org/abs/1610.02415) paper.
51+
52+
~250,000 molecular graphs with a penalized logP regression target
53+
(`y = logP - SAS - cycles`). The 12k benchmark subset follows the
54+
[Benchmarking GNNs](https://arxiv.org/abs/2003.00982) paper split.
55+
56+
# Arguments
57+
- `split` : `:train`, `:val`, or `:test`. (default: `:train`)
58+
- `subset` : Load the 12k benchmark subset instead of the full 250k. (default: `false`)
59+
- `dir` : Custom data directory. Uses DataDeps default if not set.
60+
61+
# Features
62+
- Node features: atom type, integer in `[1, 28]` → `g.node_data.features`
63+
- Edge features: bond type, integer in `[1, 4]` → `g.edge_data.bond_type`
64+
- Graph target: penalized logP, Float32 → `d.graph_data.targets[i]`
65+
66+
# Examples
67+
68+
```julia
69+
data = ZINC(split=:train, subset=true)
70+
g, y = data[1]
71+
println(g.num_nodes) # number of atoms
72+
println(g.node_data.features) # atom types (1-indexed)
73+
println(g.edge_data.bond_type) # bond types (1-indexed)
74+
println(y) # penalized logP
75+
76+
graphs, targets = data[:] # all at once
77+
batch = data[1:32] # minibatch
78+
```
79+
80+
# Dataset statistics
81+
82+
Full variant:
83+
84+
| Split | Graphs | Avg nodes | Avg edges |
85+
|-------|---------|-----------|-----------|
86+
| train | 220,011 | ~23.2 | ~49.8 |
87+
| val | 24,445 | ~23.2 | ~49.8 |
88+
| test | 5,000 | ~23.2 | ~49.8 |
89+
90+
Subset (benchmark) variant:
91+
92+
| Split | Graphs |
93+
|-------|--------|
94+
| train | 10,000 |
95+
| val | 1,000 |
96+
| test | 1,000 |
97+
"""
98+
struct ZINC <: AbstractDataset
99+
metadata::Dict{String, Any}
100+
graphs::Vector{Graph}
101+
graph_data::NamedTuple
102+
end
103+
104+
105+
function ZINC(; split = :train, subset = false, dir = nothing)
106+
@assert split in [:train, :val, :test] "split must be :train, :val, or :test"
107+
108+
root = isnothing(dir) ? datadep"ZINC" : dir
109+
split_str = String(split)
110+
data_dir = _zinc_data_dir(root)
111+
112+
npz_path = joinpath(data_dir, "$split_str.npz")
113+
isfile(npz_path) || error("Cannot find $npz_path. Ensure the ZINC data downloaded correctly.")
114+
115+
data = NPZ.npzread(npz_path)
116+
117+
# everything is stored as flat 1D arrays; node_counts/edge_counts tell us
118+
# where each molecule starts and ends
119+
atom_types_flat = data["atom_types"]
120+
edge_src_flat = data["edge_src"]
121+
edge_dst_flat = data["edge_dst"]
122+
bond_attrs_flat = data["bond_attrs"]
123+
node_counts = data["node_counts"]
124+
edge_counts = data["edge_counts"]
125+
targets_all = data["targets"]
126+
127+
node_offsets = vcat(0, cumsum(Int.(node_counts)))
128+
edge_offsets = vcat(0, cumsum(Int.(edge_counts)))
129+
130+
# for subset=true, read the official benchmark index file (0-based → 1-based)
131+
indices = if subset
132+
index_path = joinpath(root, "$split_str.index")
133+
isfile(index_path) || (index_path = joinpath(data_dir, "$split_str.index"))
134+
_read_index_file(index_path)
135+
else
136+
collect(1:length(targets_all))
137+
end
138+
139+
graphs = Vector{Graph}(undef, length(indices))
140+
targets = Vector{Float32}(undef, length(indices))
141+
142+
for (out_idx, mol_idx) in enumerate(indices)
143+
ns = node_offsets[mol_idx] + 1; ne = node_offsets[mol_idx + 1]
144+
es = edge_offsets[mol_idx] + 1; ee = edge_offsets[mol_idx + 1]
145+
146+
# Python uses 0-based atom/bond indices, Julia uses 1-based
147+
atom_type = Int.(atom_types_flat[ns:ne]) .+ 1
148+
src = Int.(edge_src_flat[es:ee]) .+ 1
149+
dst = Int.(edge_dst_flat[es:ee]) .+ 1
150+
bond_attrs = Int.(bond_attrs_flat[es:ee])
151+
152+
graphs[out_idx] = Graph(
153+
num_nodes = length(atom_type),
154+
edge_index = (src, dst),
155+
node_data = (features = atom_type,),
156+
edge_data = (bond_type = bond_attrs,),
157+
)
158+
159+
targets[out_idx] = targets_all[mol_idx]
160+
end
161+
162+
metadata = Dict{String, Any}(
163+
"split" => split,
164+
"subset" => subset,
165+
"variant" => subset ? "subset" : "full",
166+
"num_graphs" => length(graphs),
167+
"num_atom_types" => 28,
168+
"num_bond_types" => 4,
169+
"task" => "graph regression",
170+
"target" => "penalized logP (logP - SAS - cycles)",
171+
)
172+
173+
return ZINC(metadata, graphs, (targets = targets,))
174+
end
175+
176+
177+
Base.length(d::ZINC) = length(d.graphs)
178+
179+
function Base.getindex(d::ZINC, ::Colon)
180+
return (; d.graphs, d.graph_data.targets)
181+
end
182+
183+
function Base.getindex(d::ZINC, i)
184+
return getobs((; d.graphs, d.graph_data.targets), i)
185+
end
186+
187+
function Base.show(io::IO, ::MIME"text/plain", d::ZINC)
188+
recur_io = IOContext(io, :compact => false)
189+
print(io, "ZINC $(d.metadata["variant"]) - $(d.metadata["split"]):")
190+
for f in fieldnames(ZINC)
191+
startswith(string(f), "_") && continue
192+
print(recur_io, "\n $(leftalign(string(f), 12)) => $(_summary(getfield(d, f)))")
193+
end
194+
end
195+
196+
197+
# finds where the npz files live — handles fresh DataDeps download (root)
198+
# and local dev with original pickle layout (molecules/ subfolder)
199+
function _zinc_data_dir(root::String)
200+
isfile(joinpath(root, "train.npz")) && return root
201+
candidate = joinpath(root, "molecules")
202+
isdir(candidate) && return candidate
203+
error("Cannot locate ZINC data under $root.")
204+
end
205+
206+
# index files are comma-separated 0-based ints, convert to 1-based for Julia
207+
function _read_index_file(path::String)::Vector{Int}
208+
return parse.(Int, split(strip(read(path, String)), ",")) .+ 1
209+
end

‎test/datasets/graphs.jl‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,18 @@ end
175175
@test maximum(a) == g.num_nodes
176176
end
177177
end
178+
@testset "ZINC" begin
179+
data = ZINC(split=:train, subset=true)
180+
@test data isa AbstractDataset
181+
@test length(data) == 10000
182+
g, y = data[1]
183+
@test g isa MLDatasets.Graph
184+
@test g.num_nodes > 0
185+
@test g.num_edges > 0
186+
@test y isa Float32
187+
@test all(1 .<= g.node_data.features .<= 28)
188+
@test all(1 .<= g.edge_data.bond_type .<= 4)
189+
end
178190
# maybe, maybe, maybe??
179191
Sys.iswindows() || @testset "OGBn-mag" begin
180192
data = OGBDataset("ogbn-mag")

0 commit comments

Comments
 (0)