Deterministic ClusterMAFs

[yallup]

I've been using ClusterMAF for a LFT example, when trying to sample from the resulting clusterMAF as a prior in polychord (with MPI) I get a lot of non-deterministic likelihoods. This works fine when using a single process so I assume that

a. The Clustering is rebuilt on a call of flow = MAF.load('phi4/low_res_maf.pkl') b. This is not exactly deterministic if I call the same load on multiple processes?

I can get by with a single thread for now but wanted to see if this matches what would be expected with the clusterMAF as is

[htjb]

Hi @yallup, it does look like there is an issue here. I wrote the following toy script...

import numpy as np
import matplotlib.pyplot as plt
from margarine.clustered import clusterMAF

x = np.hstack([np.random.normal(0, 1, 1000), np.random.normal(2, 1, 1000)])
y = np.hstack([np.random.normal(0, 1, 1000), np.random.normal(10, 1, 1000)])
data = np.vstack([x, y]).T
print(data.shape)

flow = clusterMAF(data)
flow.train(epochs=1000, early_stop=True)
flow.save('cluster_test.pkl')

u = np.random.uniform(0, 1, (1, 2))

for i in range(10):
    flow = clusterMAF.load('cluster_test.pkl')
    print(flow(u).numpy()[0])

for i in range(10):
    print(flow(u).numpy()[0])

In the first loop I get two distinct sets of samples for the same pair of random variables u whereas in the second I get only one set of samples.

[htjb]

Ahh I think I realised partially what is happening here.

There is a random choice made when drawing samples from the clusterMAF which determines which cluster flow to draw the samples from. I think the fix might be as simple as setting the numpy random seed when loading the MAF.

In my toy example the two sets of samples I was getting when running over the first for loop were very clearly from the two different clusters.

[htjb]

@yallup can you check whether the fix on the deterministic_sampling branch fixes your non-deterministic likelihood issue please?

I just set the numpy random seed in the call function. Bit unclear to me why this is an issue in the first for loop in the toy example and not the second but this seems to resolve it. Something to do with the way the class is initialised in the load function and the point at which numpy is imported I guess.

[yallup]

Tested and looks good to me!

EDIT: Spoke too soon perhaps, the process now runs but it looks like the MAF prior collapses into just one mode, quite confusing behaviour

[htjb]

Ahh this is upsetting... I'll continue to have a play with my toy example...

[htjb]

@yallup my bad I was setting the random seed in the call function and (after some offline discussion with @ThomasGesseyJones) think it should be set in the load function. Can you try again?

[htjb]

Hmm not sure that is working... the problem here is that np.random.choice always gives the same cluster if you set the seed regardless of the samples u that are given to the call function. and if you don't set the seed then np.random.choice can give different clusters even if u is the same on repeated calls making it non deterministic.

[htjb]

@yallup I think I have it fixed now. Set the random seed for the choice of cluster and then reset it to numpy default after. Appears to work when I run the following...

EDIT: setting the random seed based on the value passed from the hypercube. This super imposes a uniform distribution over the choice of clusters but leads to deterministic sampling on multiple nodes I believe.

import numpy as np
import matplotlib.pyplot as plt
from margarine.clustered import clusterMAF

#np.random.seed(1420)

x = np.hstack([np.random.normal(0, 1, 1000), np.random.normal(2, 1, 1000)])
y = np.hstack([np.random.normal(0, 1, 1000), np.random.normal(10, 1, 1000)])
data = np.vstack([x, y]).T
print(data.shape)

flow = clusterMAF(data)
flow.train(epochs=1000, early_stop=True)
flow.save('cluster_test.pkl')

# needs to give the same result
u = np.random.uniform(0, 1, (3, 2))
for i in range(5):
    print('-----')
    # equivalent to loading the flow on 5 different nodes
    # and transforming the same samples from the unit hypercube
    # on each
    flow = clusterMAF.load('cluster_test.pkl')
    for j in range(3):
        print(flow(u[j][np.newaxis, :]).numpy()[0])

print('=====')
#needs to give the same result
flow = clusterMAF.load('cluster_test.pkl')
for i in range(5):
    # load the flow once then evaluste the
    # flow on the same samples from the hypercuve
    # five times
    print(flow(u[0][np.newaxis, :]).numpy()[0])

print('=====')
# should give different results...
for i in range(5):
    # load the flow on five different nodes
    # and evaluate on five different samples
    # from the hypercube
    flow = clusterMAF.load('cluster_test.pkl')
    u = np.random.uniform(0, 1, (1, 2))
    print(flow(u).numpy()[0])

[yallup]

Pulled the latest changes and it all looks to be passing the eye test now

[htjb]

Neat! I tried a little toy clustering example to check I hadn't broken anything. All go ahead and merge.

import numpy as np
import matplotlib.pyplot as plt
from margarine.clustered import clusterMAF
from anesthetic import MCMCSamples

x = np.hstack([np.random.normal(0, 1, 10000), np.random.normal(2, 1, 2000)])
y = np.hstack([np.random.normal(0, 1, 10000), np.random.normal(10, 1, 2000)])
data = np.vstack([x, y]).T

flow = clusterMAF(data)
flow.train(epochs=10000, early_stop=True)
flow.save('cluster_test.pkl')

s = flow.sample(10000)

data = MCMCSamples(data, columns=[0, 1])
s = MCMCSamples(s, columns=[0, 1])

axes = data.plot_2d([0, 1])
s.plot_2d(axes)
plt.savefig('cluster_test.png')
plt.show()

[htjb]

Closed with #46