def test_df():
m = Message(tensors, vectors)
df = m.dataframe()
assert df.equals(pd.DataFrame(vectors))
df = m.dataframe(keys=['c'])
assert df.equals(pd.DataFrame({'c': vectors['c']}))
df = m.dataframe(keys=['c','a'])
assert (df == (pd.DataFrame({'c': vectors['c'], 'a': np.array(tensors['a'])}))).all().all()
def test_Message_set_get():
# Test point updates
email = Message(tensors, vectors)
gmail = Message({'a':torch.Tensor([1,42,3]), 'b':torch.Tensor([4,43,6]), 'c': np.array([7,99,9]), 'd': np.array([10,100,12])})
replacement = {'a': torch.Tensor([42]), 'b': torch.Tensor([43]), 'c': np.array([99]), 'd': np.array([100])}
assert len(email) == 3
assert email != gmail
email[1] = replacement
assert email == gmail
# Test ranged updates
email = Message(tensors, vectors)
gmail = Message({'a':torch.Tensor([1,42,33]), 'b':torch.Tensor([4,43,66]), 'c': np.array([7,99,99]), 'd': np.array([10,100,122])})
replacement = {'a': torch.Tensor([42,33]), 'b': torch.Tensor([43,66]), 'c': np.array([99,99]), 'd': np.array([100,122])}
assert email != gmail
email[1:3] = replacement
assert email == gmail
# Test updates using lists as indexes
email = Message(tensors, vectors)
assert email != gmail
email[[1,2]] = replacement
assert email == gmail
# Test column updates
email['a'] = torch.Tensor([9,9,9])
assert torch.equal(email['a'], torch.Tensor([9,9,9]))
email['c'] = np.array([9,9,9])
assert email['c'].equals(pd.Series([9,9,9]))
# Test column updates that switch from df to tensor and vice-versa
email = Message(tensors, vectors)
assert set(email.columns) == set(['a','b','c','d'])
assert set(email.tensor_message.columns) == set(['a','b'])
assert set(email.df.columns) == set(['c','d'])
new_a = np.array([1,2,3]) # Switch from tensor to vector
email['a'] = new_a
assert set(email.columns) == set(['a','b','c','d'])
assert set(email.tensor_message.columns) == set(['b'])
assert set(email.df.columns) == set(['a','c','d'])
assert (email['a'] == new_a).all()
new_c = torch.Tensor([7,8,9])
email['c'] = new_c
assert set(email.columns) == set(['a','b','c','d'])
assert set(email.tensor_message.columns) == set(['b','c'])
assert set(email.df.columns) == set(['a','d'])
assert (email['c'] == new_c).all()
# Test column updates that end up clearing either self.df or self.tensor_message
email = Message(tensors, vectors)
df = email.dataframe(['a', 'b'])
assert len(email) == 3
assert len(email.tensor_message) == 3
assert len(email.df) == 3
email[['a','b']] = df
assert len(email) == 3
assert len(email.tensor_message) == 0
assert len(email.df) == 3