<filename>test/test_pool_pooling.py<gh_stars>10-100
import unittest

import numpy as np
import tensorflow as tf

from kgcnn.layers.pooling import PoolingLocalEdgesLSTM
from kgcnn.layers.gather import GatherNodes
from kgcnn.layers.modules import LazyConcatenate


class TestPoolingLocalEdgesLSTM(unittest.TestCase):

    n1 = [[[1.0], [6.0], [1.0], [6.0], [1.0], [1.0], [6.0], [6.0]],
          [[6.0], [1.0], [1.0], [1.0], [7.0], [1.0], [6.0], [8.0], [6.0], [1.0], [6.0], [7.0], [1.0], [1.0], [1.0]]]
    ei1 = [[[0, 1], [1, 0], [1, 6], [2, 3], [3, 2], [3, 5], [3, 7], [4, 7], [5, 3], [6, 1], [6, 7], [7, 3], [7, 4],
            [7, 6]],
           [[0, 6], [0, 8], [0, 9], [1, 11], [2, 4], [3, 4], [4, 2], [4, 3], [4, 6], [5, 10], [6, 0], [6, 4], [6, 14],
            [7, 8], [8, 0], [8, 7], [8, 11], [9, 0], [10, 5], [10, 11], [10, 12], [10, 13], [11, 1], [11, 8], [11, 10],
            [12, 10], [13, 10], [14, 6]]]
    e1 = [[[0.408248290463863], [0.408248290463863], [0.3333333333333334], [0.35355339059327373], [0.35355339059327373],
           [0.35355339059327373], [0.25], [0.35355339059327373], [0.35355339059327373], [0.3333333333333334],
           [0.2886751345948129], [0.25], [0.35355339059327373], [0.2886751345948129]],
          [[0.25], [0.25], [0.35355339059327373], [0.35355339059327373], [0.35355339059327373], [0.35355339059327373],
           [0.35355339059327373], [0.35355339059327373], [0.25], [0.3162277660168379], [0.25], [0.25],
           [0.35355339059327373], [0.35355339059327373], [0.25], [0.35355339059327373], [0.25], [0.35355339059327373],
           [0.3162277660168379], [0.22360679774997896], [0.3162277660168379], [0.3162277660168379],
           [0.35355339059327373], [0.25], [0.22360679774997896], [0.3162277660168379], [0.3162277660168379],
           [0.35355339059327373]]]

    def test_pooling_LSTM(self):
        n = tf.ragged.constant(self.n1, ragged_rank=1, inner_shape=(1,))
        edi = tf.ragged.constant(self.ei1, ragged_rank=1, inner_shape=(2,))
        ed = tf.ragged.constant(self.e1, ragged_rank=1, inner_shape=(1,))

        ns = GatherNodes()([n, edi])
        messages = LazyConcatenate(axis=-1)([ed, ns])
        out = PoolingLocalEdgesLSTM(units=3)([n, messages, edi])
        print(out[0].shape)
        self.assertTrue(np.all(np.array(out[0].shape) == np.array([8,3])))


if __name__ == '__main__':
    unittest.main()



StarcoderdataPython