1 files changed, 95 insertions, 0 deletions
diff --git a/simulation/mle_blocks.py b/simulation/mle_blocks.py
new file mode 100644
index 0000000..5acebab
--- /dev/null
+++ b/simulation/mle_blocks.py
@@ -0,0 +1,95 @@
+import main as mn
+import theano
+from theano import tensor as tsr
+import blocks
+import blocks.algorithms, blocks.main_loop, blocks.extensions.monitoring
+import theano.tensor.shared_randomstreams
+import picklable_itertools
+import numpy as np
+from six.moves import range
+import fuel
+import fuel.datasets
+import collections
+
+
+class JeaninuScheme(fuel.schemes.ShuffledScheme):
+    def get_request_iterator(self):
+        indices = list(self.indices)
+        start = np.random.randint(self.batch_size)
+        batches = list(map(
+            list,
+            picklable_itertools.extras.partition_all(self.batch_size,
+                                                     indices[start:])
+        ))
+        if indices[:start]:
+            batches.append(indices[:start])
+        batches = np.asarray(batches)
+        return iter(batches[np.random.permutation(len(batches))])
+
+
+def create_model(n_nodes):
+    x = tsr.matrix(name='x', dtype='int8')
+    s = tsr.matrix(name='s', dtype='int8')
+    params = theano.shared(
+        .5 + .01 *
+        np.random.normal(size=(n_nodes, n_nodes)).astype(theano.config.floatX),
+        name='params'
+    )
+    y = tsr.maximum(tsr.dot(x, params), 1e-5)
+    infect = tsr.log(1. - tsr.exp(-y[0:-1]))
+    lkl_pos = tsr.sum(infect * (x[1:] & s[1:]))
+    lkl_neg = tsr.sum(-y[0:-1] * (~x[1:] & s[1:]))
+    lkl_mle = lkl_pos + lkl_neg
+    lkl_mle.name = 'cost'
+    return x, s, params, lkl_mle
+
+
+def create_random_graph(n_nodes, p=.5):
+    graph = .5 * np.random.binomial(2, p=.5, size=(n_nodes, n_nodes))
+    for k in range(len(graph)):
+        graph[k, k] = 0
+    return np.log(1. / (1 - p * graph))
+
+
+def create_data_stream(n_cascades, graph, batch_size, shuffle=True):
+    """
+    shuffle (bool): shuffle minibatches but not within minibatch
+    """
+    cascades = mn.build_cascade_list(mn.simulate_cascades(n_cascades, graph),
+                                     collapse=True)
+    x_obs, s_obs = cascades[0], cascades[1]
+    data_set = fuel.datasets.base.IndexableDataset(collections.OrderedDict(
+        [('x', x_obs), ('s', s_obs)]
+    ))
+    if shuffle:
+        scheme = JeaninuScheme(len(x_obs), batch_size=batch_size)
+    else:
+        scheme = fuel.schemes.SequentialScheme(len(x_obs),
+                                               batch_size=batch_size)
+    return fuel.streams.DataStream(dataset=data_set, iteration_scheme=scheme)
+
+
+if __name__ == "__main__":
+    n_cascades = 10000
+    batch_size = 1000
+    graph = np.array([[0, 0, 1], [0, 0, 0.5], [0, 0, 0]])
+    graph = np.log(1. / (1 - .5 * graph))
+    print('GRAPH:\n', graph, '\n-------------\n')
+
+    x, s, params, cost = create_model(len(graph))
+
+    alg = blocks.algorithms.GradientDescent(
+       cost=-cost, parameters=[params], step_rule=blocks.algorithms.AdaDelta()
+    )
+    data_stream = create_data_stream(n_cascades, graph, batch_size,
+            shuffle=True)
+    loop = blocks.main_loop.MainLoop(
+        alg, data_stream,
+        extensions=[
+            blocks.extensions.FinishAfter(after_n_epochs = 1000),
+            blocks.extensions.monitoring.TrainingDataMonitoring([cost, params],
+                after_batch=True),
+            blocks.extensions.Printing()
+        ]
+    )
+    loop.run()