def toSparseTensor(X): X = X.tocoo() indices = np.vstack((X.row, X.col)).T # Transpose values = X.data shape = X.shape return SparseTensor(indices, values, shape)
def generate_edge_features(self, x, a): send = a.indices[:, 0] receive = a.indices[:, 1] if self.forward == True: forwards = tf.gather(x[:, 3], send) <= tf.gather(x[:, 3], receive) send = tf.cast(send[forwards], tf.int64) receive = tf.cast(receive[forwards], tf.int64) a = SparseTensor(indices=tf.stack([send, receive], axis=1), values=tf.ones(tf.shape(send), dtype=tf.float32), dense_shape=tf.cast(tf.shape(a), tf.int64)) diff_x = tf.subtract(tf.gather(x, receive), tf.gather(x, send)) dists = tf.sqrt(tf.reduce_sum(tf.square(diff_x[:, :3]), axis=1)) vects = tf.math.divide_no_nan(diff_x[:, :3], tf.expand_dims(dists, axis=-1)) e = tf.concat( [diff_x[:, 3:], tf.expand_dims(dists, -1), vects], axis=1) return a, e
def generate_edge_features(self, x, a): send = a.indices[:, 0] receive = a.indices[:, 1] if self.forward == True: #could maybe be improved forwards = tf.gather(x[:, 3], send) <= tf.gather(x[:, 3], receive) send = tf.cast(send[forwards], tf.int64) receive = tf.cast(receive[forwards], tf.int64) a = SparseTensor(indices = tf.stack([send, receive], axis = 1), values = tf.ones(tf.shape(send), dtype = tf.float32), dense_shape = tf.cast(tf.shape(a), tf.int64)) ##distance vectors diff_x = tf.subtract(tf.gather(x, receive), tf.gather(x, send)) dists = tf.sqrt( tf.reduce_sum( tf.square( diff_x[:, :3] ), axis = 1 )) vects = tf.math.divide_no_nan(diff_x[:, :3], tf.expand_dims(dists, axis = -1)) if self.edgetype==0: e = tf.concat([diff_x[:, 3:], tf.expand_dims(dists, -1), vects], axis = 1) if self.edgetype==1: ## SRT, could make this is a mask prod_x = tf.math.multiply(tf.gather(x, receive), tf.gather(x, send)) srt = prod_x[:,5] e = tf.concat([diff_x[:, 3:], tf.expand_dims(dists, -1), vects, tf.expand_dims(srt, -1)], axis = 1) if self.edgetype==2: st=2699 # time scale, database specific c=tf.constant(0.000299792458) #speed of light in km pr nanosec speed = tf.math.divide_no_nan(dists, st*diff_x[:,3]) #could add fudge factor to account for ice c lower than vacuum c speed = tf.math.greater_equal(speed, c) speed = tf.cast(tf.where(speed, 0,1), tf.float32) e = tf.concat([diff_x[:, 3:], tf.expand_dims(dists, -1), vects, tf.expand_dims(speed,-1)], axis = 1) if self.edgetype==3: #srt comm? prod_x = tf.math.multiply(tf.gather(x, receive), tf.gather(x, send)) srt = prod_x[:,5] #higher than c? st=2699 # time scale, database specific c=tf.constant(0.000299792458) #speed of light in km pr nanosec speed = tf.math.divide_no_nan(dists, st*diff_x[:,3]) #could add fudge factor to account for ice c lower than vacuum c speed = tf.math.greater_equal(speed, c) speed = tf.cast(tf.where(speed, 0,1), tf.float32) e = tf.concat([diff_x[:, 3:], tf.expand_dims(dists, -1), vects, tf.expand_dims(srt, -1), tf.expand_dims(speed,-1)], axis = 1) return a, e