How to use the elfi.tools.vectorize function in elfi

def test_dict_output():
    vsim = elfi.tools.vectorize(simulator)
    vsum = elfi.tools.vectorize(summary)

    obs = simulator([.2, .8])

    elfi.new_model()
    p = elfi.Prior('dirichlet', [2, 2])
    sim = elfi.Simulator(vsim, p, observed=obs)
    S = elfi.Summary(vsum, sim)
    d = elfi.Distance('euclidean', S)

    pool = elfi.OutputPool(['sim'])
    rej = elfi.Rejection(d, batch_size=100, pool=pool, output_names=['sim'])
    sample = rej.sample(100, n_sim=1000)
    mean = np.mean(sample.samples['p'], axis=0)

    # Crude test
    assert mean[1] > mean[0]

def test_as_vectorized_summary(self):
        ret1 = np.array([5])
        ret2 = np.array([6])
        mock_seq = elfi.tools.vectorize(MockSequentialSummary([ret1, ret2]))
        mock_vec = MockSummary([ret1, ret2])
        input_data = [np.atleast_2d([[1], [2]]), np.atleast_2d([[3], [4]])]
        output_seq = mock_seq(*input_data)
        output_vec = mock_vec(*input_data)
        assert np.array_equal(output_seq, output_vec)

def test_list_output():
    vsim = elfi.tools.vectorize(lsimulator)
    vsum = elfi.tools.vectorize(lsummary)

    v = vsim(np.array([[.2, .8], [.3, .7]]))
    assert is_array(v)
    assert not isinstance(v[0], list)

    vsim = elfi.tools.vectorize(lsimulator, dtype=False)

    v = vsim(np.array([[.2, .8], [.3, .7]]))
    assert is_array(v)
    assert isinstance(v[0], list)

    obs = lsimulator([.2, .8])

    elfi.new_model()
    p = elfi.Prior('dirichlet', [2, 2])
    sim = elfi.Simulator(vsim, p, observed=obs)

    @elfi.tools.vectorize
    def mock_simulator1(*data, random_state=None):
        return np.array((1,))

def test_as_vectorized_discrepancy(self):
        ret1 = np.array([5])
        ret2 = np.array([6])
        mock_seq = elfi.tools.vectorize(MockSequentialDiscrepancy([ret1, ret2]))
        mock_vec = MockDiscrepancy([ret1, ret2])
        x = (np.atleast_2d([[1], [2]]), np.atleast_2d([[3], [4]]))
        y = (np.atleast_2d([[5]]), np.atleast_2d([[6]]))
        output_seq = mock_seq(x, y)
        output_vec = mock_vec(x, y)
        assert np.array_equal(output_seq, output_vec)

def test_vectorized_and_external_combined():
    constant = elfi.Constant(123)
    kwargs_sim = elfi.tools.external_operation(
        'echo {seed} {batch_index} {index_in_batch} {submission_index}', process_result='int32')
    kwargs_sim = elfi.tools.vectorize(kwargs_sim)
    sim = elfi.Simulator(kwargs_sim, constant)

    with pytest.raises(Exception):
        sim.generate(3)

    sim['_uses_meta'] = True
    g = sim.generate(3)

    # Test uniqueness of seeds
    assert len(np.unique(g[:, 0]) == 3)

    assert len(np.unique(g[:, 1]) == 1)

    # Test index_in_batch
    assert np.array_equal(g[:, 2], [0, 1, 2])

def test_vectorize_decorator():
    batch_size = 3
    a = np.array([1, 2, 3])
    b = np.array([3, 2, 1])

    def simulator(a, b, random_state=None):
        return a * b

    vsim = elfi.tools.vectorize(simulator)
    assert np.array_equal(a * b, vsim(a, b, batch_size=batch_size))

    def simulator(a, constant, random_state=None):
        return a * constant

    vsim = elfi.tools.vectorize(simulator, constants=[1])
    assert np.array_equal(a * 5, vsim(a, 5, batch_size=batch_size))

    vsim = elfi.tools.vectorize(simulator, [1])
    assert np.array_equal(a * 5, vsim(a, 5, batch_size=batch_size))

    def simulator(constant0, b, constant2, random_state=None):
        return constant0 * b * constant2

    vsim = elfi.tools.vectorize(simulator, constants=(0, 2))
    assert np.array_equal(2 * b * 7, vsim(2, b, 7, batch_size=batch_size))

def test_list_output():
    vsim = elfi.tools.vectorize(lsimulator)
    vsum = elfi.tools.vectorize(lsummary)

    v = vsim(np.array([[.2, .8], [.3, .7]]))
    assert is_array(v)
    assert not isinstance(v[0], list)

    vsim = elfi.tools.vectorize(lsimulator, dtype=False)

    v = vsim(np.array([[.2, .8], [.3, .7]]))
    assert is_array(v)
    assert isinstance(v[0], list)

    obs = lsimulator([.2, .8])

    elfi.new_model()
    p = elfi.Prior('dirichlet', [2, 2])

    @elfi.tools.vectorize
    def mock_discrepancy1(x, y):
        return np.ones((1,))

def test_dict_output():
    vsim = elfi.tools.vectorize(simulator)
    vsum = elfi.tools.vectorize(summary)

    obs = simulator([.2, .8])

    elfi.new_model()
    p = elfi.Prior('dirichlet', [2, 2])
    sim = elfi.Simulator(vsim, p, observed=obs)
    S = elfi.Summary(vsum, sim)
    d = elfi.Distance('euclidean', S)

    pool = elfi.OutputPool(['sim'])
    rej = elfi.Rejection(d, batch_size=100, pool=pool, output_names=['sim'])
    sample = rej.sample(100, n_sim=1000)
    mean = np.mean(sample.samples['p'], axis=0)

    # Crude test