def test_join():
t = tensors
v = vectors
t2 = {'d': torch.Tensor([13,14,15])}
v2 = {'e': np.array([16,17,18])}
m1 = Message(t,v)
m2 = Message(t)
m2_t = TensorMessage(t)
m3 = Message(v)
m4 = Message(t2,v2)
m5 = Message(t2)
m5_t = TensorMessage(t2)
m6 = Message(v2)
m7 = Message(t,v2)
m8 = Message(t2, v)
# Test if a tensor message can be merged into a message and vice versa
assert m2.merge(m3) == m1
assert m3.merge(m2) == m1
assert m3.merge(m2_t) == m1
assert m3.merge(t) == m1
# Test if the tensors in messages can be merged
assert m2.merge(t2) == Message({**t, **t2})
assert m2.merge(m5) == Message({**t, **t2})
assert m2.merge(m5_t) == Message({**t, **t2})
assert m2_t.merge(t2) == TensorMessage({**t, **t2})
assert m2_t.merge(m5) == TensorMessage({**t, **t2})
assert m2_t.merge(m5_t) == TensorMessage({**t, **t2})
# Test if the dataframes in messages can be merged
assert m3.merge(m6) == Message({**v, **v2})
assert m6.merge(m3) == Message({**v, **v2})
assert m3.merge(v2) == Message({**v, **v2})
def drop_df(batch, exceptions=['SampleID']):
"""
Drops the dataframe component and leaves only the tensor component of a message.
"""
keep_df = batch[exceptions]
new_batch = Message(keep_df)
new_batch = new_batch.merge(batch.tensors())
return new_batch