File size: 4,310 Bytes
7180154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | # Copyright 2024 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS-IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A Transformer model for weather predictions.
This model wraps the a transformer model and swaps the leading two axes of the
nodes in the input graph prior to evaluating the model to make it compatible
with a [nodes, batch, ...] ordering of the inputs.
"""
from typing import Any, Mapping, Optional
from . import typed_graph
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from scipy import sparse
Kwargs = Mapping[str, Any]
def _get_adj_matrix_for_edge_set(
graph: typed_graph.TypedGraph,
edge_set_name: str,
add_self_edges: bool,
):
"""Returns the adjacency matrix for the given graph and edge set."""
# Get nodes and edges of the graph.
edge_set_key = graph.edge_key_by_name(edge_set_name)
sender_node_set, receiver_node_set = edge_set_key.node_sets
# Compute number of sender and receiver nodes.
sender_n_node = graph.nodes[sender_node_set].n_node[0]
receiver_n_node = graph.nodes[receiver_node_set].n_node[0]
# Build adjacency matrix.
adj_mat = sparse.csr_matrix((sender_n_node, receiver_n_node), dtype=np.bool_)
edge_set = graph.edges[edge_set_key]
s, r = edge_set.indices
adj_mat[s, r] = True
if add_self_edges:
# Should only do this if we are certain the adjacency matrix is square.
assert sender_node_set == receiver_node_set
adj_mat[np.arange(sender_n_node), np.arange(receiver_n_node)] = True
return adj_mat
class MeshTransformer(hk.Module):
"""A Transformer for inputs with ordering [nodes, batch, ...]."""
def __init__(self,
transformer_ctor,
transformer_kwargs: Kwargs,
name: Optional[str] = None):
"""Initialises the Transformer model.
Args:
transformer_ctor: Constructor for transformer.
transformer_kwargs: Kwargs to pass to the transformer module.
name: Optional name for haiku module.
"""
super().__init__(name=name)
# We defer the transformer initialisation to the first call to __call__,
# where we can build the mask senders and receivers of the TypedGraph
self._batch_first_transformer = None
self._transformer_ctor = transformer_ctor
self._transformer_kwargs = transformer_kwargs
@hk.name_like('__init__')
def _maybe_init_batch_first_transformer(self, x: typed_graph.TypedGraph):
if self._batch_first_transformer is not None:
return
self._batch_first_transformer = self._transformer_ctor(
adj_mat=_get_adj_matrix_for_edge_set(
graph=x,
edge_set_name='mesh',
add_self_edges=True,
),
**self._transformer_kwargs,
)
def __call__(
self, x: typed_graph.TypedGraph,
global_norm_conditioning: jax.Array
) -> typed_graph.TypedGraph:
"""Applies the model to the input graph and returns graph of same shape."""
if set(x.nodes.keys()) != {'mesh_nodes'}:
raise ValueError(
f'Expected x.nodes to have key `mesh_nodes`, got {x.nodes.keys()}.'
)
features = x.nodes['mesh_nodes'].features
if features.ndim != 3: # pytype: disable=attribute-error # jax-ndarray
raise ValueError(
'Expected `x.nodes["mesh_nodes"].features` to be 3, got'
f' {features.ndim}.'
) # pytype: disable=attribute-error # jax-ndarray
# Initialise transformer and mask.
self._maybe_init_batch_first_transformer(x)
y = jnp.transpose(features, axes=[1, 0, 2])
y = self._batch_first_transformer(y, global_norm_conditioning)
y = jnp.transpose(y, axes=[1, 0, 2])
x = x._replace(
nodes={
'mesh_nodes': x.nodes['mesh_nodes']._replace(
features=y.astype(features.dtype)
)
}
)
return x
|