CW nets template optimization

[erhc]

Hello,

I am trying to implement CW networks from the following paper.
I think conceptually I have it somehow written down in the form of a template, but the problem is that it takes approximately 20 minutes to ground on Mutagenesis dataset. I would appreciate if you could help me identify places that could be optimized, but any other feedback would be much appreciated!
Thank you!

Here is the template:

template = Template()

template.add_rules([
    (R.atom_embed(V.A)[3,] <= R.get(atom)(V.A)) for atom in ["c", "o", "br", "i", "f", "h", "n", "cl"]
])


template.add_rules([
    (R.bond_embed(V.B)[3,] <= R.get(bond)(V.B)) for bond in ["b_1", "b_2", "b_3", "b_4", "b_5", "b_7"]
])

# Aggregating messages bottom to top
# atoms aggregate to bonds, bonds to rings
template += R.bond_features(V.B) <= (R.bond(V.X, V.Y, V.B), R.atom_embed(V.X)[3, 3], R.atom_embed(V.Y)[3, 3])

# bonds in same cycle
def get_bond_cycle(n):
    body = [R.bond(f"X{i}", f"X{(i+1) % n}", f"B{i}") for i in range(n)]
    body.extend(R.bond_features(f"B{i}")[3, 3] for i in range(n))
    body.append(R.special.alldiff(...))

    return R.cycle_features([f"B{i}" for i in range(n)]) <= body

def connect_cycle(n):
    return R.cycle_features([f"B{i}" for i in range(n-1)]) <= R.cycle_features([f"B{i}" for i in range(n)])

max_ring_size = 7
for i in range(3, max_ring_size):
    template += get_bond_cycle(i)
    template += connect_cycle(i)

# 
template += R.features <= (R.bond_features(V.B)[3, 3])
template += R.features <= (R.cycle_features(V.A, V.B)[3, 3])

#---------------------------------------------------------------

# Aggregating messages top to bottom

def get_cycle_messages(n):
    body = [R.bond(f"X{i}", f"X{(i+1) % n}", f"B{i}") for i in range(n)]
    body.extend(R.bond_embed(f"B{i}")[3, 3] for i in range(n))
    body.append(R.special.alldiff(...))

    return R.cycle_message(V.B0) <= body


for i in range(3, max_ring_size):
    template += get_cycle_messages(i)

# atoms sharing a bond share messages, bonds in the same ring
template += R.atom_nbhood(V.X) <= (R.bond(V.X, V.Y, V.B), R.atom_embed(V.Y)[3, 3], R.bond_embed(V.B)[3, 3])
template += R.bond_nbhood(V.B) <= (R.cycle_message(V.B)) # aggregate messages from all bonds in the same ring

template += R.nbhood <= (R.bond_nbhood(V.B)[3, 3])
template += R.nbhood <= (R.atom_nbhood(V.B)[3, 3])
#---------------------------------------------------------------
template += R.predict[1, 3] <= (R.nbhood)
template += R.predict[1, 3] <= (R.features)

[LukasZahradnik]

Hello,

It appears that the slow part are the cycle_features rules. Could you describe what the cycle_features is supposed to do?

[erhc]

Thank you for your reply, in this model there are three "layers", atoms which when connected form bonds, and which in turn form rings or cycles (something like a set of bonds enclosing a space, this is taken from topology). This predicate is meant to somehow pass messages from bonds to these cycles, just like messages are passed from atoms to bonds in bond_features predicate, and then I would connect them in a way that is agnostic to the size of the cycle.

[LukasZahradnik]

Okay, so for each cycle of a given length, you aggregate bond features in each cycle. Then you just aggregate all of those new cycle representations together?

If so, then you could simplify the template quite a bit, and instead of having the following (your current template):

cycle_features(B0, B1, B2) :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X0, B2), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), @alldiff(X1, B2, B0, X0, X2, B1).
cycle_features(B0, B1) :- cycle_features(B0, B1, B2).
cycle_features(B0, B1, B2, B3) :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X3, B2), bond(X3, X0, B3), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), {3, 3} bond_features(B3), @alldiff(B3, X1, B2, B0, X3, X0, X2, B1).
cycle_features(B0, B1, B2) :- cycle_features(B0, B1, B2, B3).
cycle_features(B0, B1, B2, B3, B4) :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X3, B2), bond(X3, X4, B3), bond(X4, X0, B4), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), {3, 3} bond_features(B3), {3, 3} bond_features(B4), @alldiff(X4, B3, X1, B2, B0, X3, X0, X2, B4, B1).
cycle_features(B0, B1, B2, B3) :- cycle_features(B0, B1, B2, B3, B4).
cycle_features(B0, B1, B2, B3, B4, B5) :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X3, B2), bond(X3, X4, B3), bond(X4, X5, B4), bond(X5, X0, B5), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), {3, 3} bond_features(B3), {3, 3} bond_features(B4), {3, 3} bond_features(B5), @alldiff(X4, B3, X5, X1, B2, B0, X3, X0, X2, B4, B5, B1).
cycle_features(B0, B1, B2, B3, B4) :- cycle_features(B0, B1, B2, B3, B4, B5).

You could have the following (which builds under 10s on my machine):

cycle_features :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X0, B2), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), @alldiff(X2, X1, B0, X0, B2, B1).
cycle_features :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X3, B2), bond(X3, X0, B3), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), {3, 3} bond_features(B3), @alldiff(X2, X1, X3, B0, X0, B2, B1, B3).
cycle_features :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X3, B2), bond(X3, X4, B3), bond(X4, X0, B4), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), {3, 3} bond_features(B3), {3, 3} bond_features(B4), @alldiff(X2, X1, X3, B0, X0, X4, B2, B1, B4, B3).
cycle_features :- bond(X0, X1, B0), bond(X1, X2, B1), bond(X2, X3, B2), bond(X3, X4, B3), bond(X4, X5, B4), bond(X5, X0, B5), {3, 3} bond_features(B0), {3, 3} bond_features(B1), {3, 3} bond_features(B2), {3, 3} bond_features(B3), {3, 3} bond_features(B4), {3, 3} bond_features(B5), @alldiff(X2, X1, X3, B0, B5, X0, X4, B2, B1, X5, B4, B3).
features :- {3, 3} bond_features(B).
features :- {3, 3} cycle_features.

Or is there a reason for passing around all those bond variables in rule heads?

[erhc]

Oh yes, this is painfully obvious now... With the bond variables I wanted to make sure the messages are passed only between the bonds in the same cycle, but I see your approach will do exactly that. It works perfectly now, thank you!