def compute_coarse_matrix_sparse(self, model, A_stencil, A_matrices, grid_size, bb=True):
if bb == True:
P_stencil = model(inputs=A_stencil, black_box=True)
else:
P_stencil = model(inputs=A_stencil, black_box=False, phase="Test")
P_matrix = tf.to_double(self.compute_p2_sparse(P_stencil, P_stencil.shape.as_list()[0], grid_size))
P_matrix_t = tf.sparse_transpose(P_matrix, [0, 2, 1])
A_matrices = tf.squeeze(A_matrices)
temp = tf.sparse_tensor_to_dense(P_matrix_t)
q = tf.matmul(temp, tf.to_double(A_matrices))
A_c = tf.transpose(tf.matmul(temp, tf.transpose(q, [0, 2, 1])), [0, 2, 1])
return A_c, self.compute_stencil(tf.squeeze(A_c), (grid_size // 2)), P_matrix, P_matrix_t
def _most_recent_obs_value(obs_values, indicator, delta_time,
attribution_max_delta_time):
"""Returns the most recent lab result for each test within a time frame.
The eligible lab values fall into a time window until time of prediction -
attribution_max_delta_time. Among those we select their most recent value
or zero if there are none.
Args:
obs_values: A dense representation of the observation_values at the position
of their obs_code_ids. A padded Tensor of shape [batch_size,
max_sequence_length, vocab_size] of type float32 where obs_values[b, t,
id] = observation_values[b, t, 0] and id = observation_code_ids[b, t, 0]
and obs_values[b, t, x] = 0 for all other x != id. If t is greater than
the sequence_length of batch entry b then the result is 0 as well.
indicator: A one-hot encoding of whether a value in obs_values comes from
observation_values or is just filled in to be 0. A Tensor of shape
[batch_size, max_sequence_length, vocab_size] and type float32.
delta_time: A Tensor of shape [batch_size, max_sequence_length] describing
the time to prediction.
attribution_max_delta_time: Time threshold so that we return the most recent
lab values among those that are at least attribution_max_delta_time
seconds old at time of prediction.
Returns:
A Tensor of shape [batch_size, 1, vocab_size] of the most recent lab results
for all lab tests that are at least attribution_max_delta_time old at time
of prediction.
"""
batch_size = tf.shape(indicator)[0]
seq_len = tf.shape(indicator)[1]
num_obs = indicator.shape[2]
# Prepend a dummy so that for lab tests for which we have no eligible lab
# values we will select 0.
obs_values = tf.concat(
[tf.zeros([batch_size, 1, num_obs]), obs_values], axis=1)
indicator = tf.concat([tf.ones([batch_size, 1, num_obs]), indicator], axis=1)
delta_time = tf.to_int32(delta_time)
delta_time = tf.concat(
[
tf.zeros([batch_size, 1, 1], dtype=tf.int32) +
attribution_max_delta_time, delta_time
],
axis=1)
# First we figure out what the eligible lab values are that are at least
# attribution_max_delta_time old.
indicator = tf.to_int32(indicator)
indicator *= tf.to_int32(delta_time >= attribution_max_delta_time)
range_val = tf.expand_dims(tf.range(seq_len + 1), axis=0)
range_val = tf.tile(range_val, multiples=[tf.shape(indicator)[0], 1])
# [[[0], [1], ..., [max_sequence_length]],
# [[0], [1], ..., [max_sequence_length]],
# ...]
range_val = tf.expand_dims(range_val, axis=2)
# [batch_size, max_sequence_length, vocab_size] with 1 non-zero number per
# time-step equal to that time-step.
seq_indicator = indicator * range_val
# [batch_size, vocab_size] with the time-step of the last lab value.
last_val_indicator = tf.reduce_max(seq_indicator, axis=1, keepdims=True)
last_val_indicator = tf.tile(
last_val_indicator, multiples=[1, tf.shape(indicator)[1], 1])
# eq indicates which lab values are the most recent ones.
eq = tf.logical_and(
tf.equal(last_val_indicator, seq_indicator), indicator > 0)
most_recent_obs_value_indicator = tf.where(eq)
# Collect the lab values associated with those indices.
res = tf.gather_nd(obs_values, most_recent_obs_value_indicator)
# Reorder the values by batch and then by lab test.
res_sorted = tf.sparse_reorder(
tf.sparse_transpose(
tf.SparseTensor(
indices=most_recent_obs_value_indicator,
values=res,
dense_shape=tf.to_int64(
tf.stack([batch_size, seq_len + 1, num_obs]))),
perm=[0, 2, 1])).values
return tf.reshape(res_sorted, [batch_size, 1, num_obs])
def sparse_message_pass(node_states,
adjacency_matrices,
num_edge_types,
hidden_size,
use_bias=True,
average_aggregation=False,
name="sparse_ggnn"):
"""One message-passing step for a GNN with a sparse adjacency matrix.
Implements equation 2 (the message passing step) in
[Li et al. 2015](https://arxiv.org/abs/1511.05493).
N = The number of nodes in each batch.
H = The size of the hidden states.
T = The number of edge types.
Args:
node_states: Initial states of each node in the graph. Shape is [N, H].
adjacency_matrices: Adjacency matrix of directed edges for each edge
type. Shape is [N, N, T] (sparse tensor).
num_edge_types: The number of edge types. T.
hidden_size: The size of the hidden state. H.
use_bias: Whether to use bias in the hidden layer.
average_aggregation: How to aggregate the incoming node messages. If
average_aggregation is true, the messages are averaged. If it is false,
they are summed.
name: (optional) The scope within which tf variables should be created.
Returns:
The result of one step of Gated Graph Neural Network (GGNN) message passing.
Shape: [N, H]
"""
n = tf.shape(node_states)[0]
t = num_edge_types
incoming_edges_per_type = tf.sparse_reduce_sum(adjacency_matrices, axis=1)
# Convert the adjacency matrix into shape [T, N, N] - one [N, N] adjacency
# matrix for each edge type. Since sparse tensor multiplication only supports
# two-dimensional tensors, we actually convert the adjacency matrix into a
# [T * N, N] tensor.
adjacency_matrices = tf.sparse_transpose(adjacency_matrices, [2, 0, 1])
adjacency_matrices = tf.sparse_reshape(adjacency_matrices, [t * n, n])
# Multiply the adjacency matrix by the node states, producing a [T * N, H]
# tensor. For each (edge type, node) pair, this tensor stores the sum of
# the hidden states of the node's neighbors over incoming edges of that type.
messages = tf.sparse_tensor_dense_matmul(adjacency_matrices, node_states)
# Rearrange this tensor to have shape [N, T * H]. The incoming states of each
# nodes neighbors are summed by edge type and then concatenated together into
# a single T * H vector.
messages = tf.reshape(messages, [t, n, hidden_size])
messages = tf.transpose(messages, [1, 0, 2])
messages = tf.reshape(messages, [n, t * hidden_size])
# Run each of those T * H vectors through a linear layer that produces
# a vector of size H. This process is equivalent to running each H-sized
# vector through a separate linear layer for each edge type and then adding
# the results together.
#
# Note that, earlier on, we added together all of the states of neighbors
# that were connected by edges of the same edge type. Since addition and
# multiplying by a linear layer are commutative, this process was equivalent
# to running each incoming edge through a linear layer separately and then
# adding everything at the end.
with tf.variable_scope(name, default_name="sparse_ggnn"):
final_node_states = common_layers.dense(messages,
hidden_size,
use_bias=False)
# Multiply the bias by for each edge type by the number of incoming nodes
# of that edge type.
if use_bias:
bias = tf.get_variable("bias",
initializer=tf.zeros([t, hidden_size]))
final_node_states += tf.matmul(incoming_edges_per_type, bias)
if average_aggregation:
incoming_edges = tf.reduce_sum(incoming_edges_per_type,
-1,
keepdims=True)
incoming_edges = tf.tile(incoming_edges, [1, hidden_size])
final_node_states /= incoming_edges + 1e-7
return tf.reshape(final_node_states, [n, hidden_size])