k3-node 1.0.0__py3-none-any.whl

k3_node/models/metapath2vec.py

@@ -0,0 +1,234 @@
+from typing import Dict, List, Optional, Tuple
+import numpy as np
+import keras
+from keras import ops
+
+EdgeType = Tuple[str, str, str]
+NodeType = str
+
+
+class MetaPath2Vec(keras.layers.Layer):
+    r"""The MetaPath2Vec model from the `"metapath2vec: Scalable Representation
+    Learning for Heterogeneous Networks"
+    <https://ericdongyx.github.io/papers/
+    KDD17-dong-chawla-swami-metapath2vec.pdf>`_ paper where random walks based
+    on a given :obj:`metapath` are sampled in a heterogeneous graph, and node
+    embeddings are learned via negative sampling optimization.
+
+    Args:
+        edge_index_dict (Dict[Tuple[str, str, str], Tensor]): Dictionary
+            holding edge indices for each edge type.
+        embedding_dim (int): The size of each embedding vector.
+        metapath (List[Tuple[str, str, str]]): The sequence of edge types
+            denoting the metapath.
+        walk_length (int): The walk length.
+        context_size (int): The context size considered for positive samples.
+        walks_per_node (int, optional): The number of walks to sample for each node.
+            (default: :obj:`1`)
+        num_negative_samples (int, optional): The number of negative samples.
+            (default: :obj:`1`)
+        num_nodes_dict (Dict[str, int], optional): The number of nodes for each
+            node type. (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MetaPath2Vec
+
+        edge_index_dict = {
+            ("author", "writes", "paper"): np.array([[0, 1, 1], [0, 0, 1]]),
+            ("paper", "written_by", "author"): np.array([[0, 0, 1], [0, 1, 1]]),
+        }
+        metapath = [("author", "writes", "paper"), ("paper", "written_by", "author")]
+        model = MetaPath2Vec(edge_index_dict, embedding_dim=16, metapath=metapath,
+                             walk_length=2, context_size=2, walks_per_node=2)
+        print(tuple(model("author").shape))  # (2, 16): embeddings of all authors
+        print(tuple(model("paper").shape))  # (2, 16)
+        ```
+    """
+    def __init__(
+        self,
+        edge_index_dict: Dict[EdgeType, any],
+        embedding_dim: int,
+        metapath: List[EdgeType],
+        walk_length: int,
+        context_size: int,
+        walks_per_node: int = 1,
+        num_negative_samples: int = 1,
+        num_nodes_dict: Optional[Dict[NodeType, int]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if num_nodes_dict is None:
+            num_nodes_dict = {}
+            for keys, edge_index in edge_index_dict.items():
+                e_np = ops.convert_to_numpy(edge_index)
+                key = keys[0]
+                N = int(e_np[0].max() + 1) if e_np.size > 0 else 0
+                num_nodes_dict[key] = max(N, num_nodes_dict.get(key, N))
+
+                key = keys[-1]
+                N = int(e_np[1].max() + 1) if e_np.size > 0 else 0
+                num_nodes_dict[key] = max(N, num_nodes_dict.get(key, N))
+
+        # Build adjacency dictionaries
+        self.adj_dict = {}
+        for keys, edge_index in edge_index_dict.items():
+            src_type, _, dst_type = keys
+            num_src = num_nodes_dict[src_type]
+            adj = [[] for _ in range(num_src)]
+            e_np = ops.convert_to_numpy(edge_index).astype(np.int64)
+            if e_np.size > 0:
+                for s, d in zip(e_np[0], e_np[1]):
+                    adj[s].append(int(d))
+            self.adj_dict[keys] = adj
+
+        for edge_type1, edge_type2 in zip(metapath[:-1], metapath[1:]):
+            if edge_type1[-1] != edge_type2[0]:
+                raise ValueError(
+                    "Found invalid metapath. Ensure that the destination node "
+                    "type matches with the source node type across all "
+                    "consecutive edge types."
+                )
+
+        assert walk_length + 1 >= context_size
+
+        self.embedding_dim = embedding_dim
+        self.metapath = metapath
+        self.walk_length = walk_length
+        self.context_size = context_size
+        self.walks_per_node = walks_per_node
+        self.num_negative_samples = num_negative_samples
+        self.num_nodes_dict = num_nodes_dict
+        self.EPS = 1e-15
+
+        types = sorted(list({x[0] for x in metapath} | {x[-1] for x in metapath}))
+
+        count = 0
+        self.start, self.end = {}, {}
+        for key in types:
+            self.start[key] = count
+            count += num_nodes_dict[key]
+            self.end[key] = count
+
+        offset = [self.start[metapath[0][0]]]
+        offset += [self.start[keys[-1]] for keys in metapath] * int(
+            (walk_length / len(metapath)) + 1
+        )
+        offset = offset[: walk_length + 1]
+        self.offset = np.array(offset, dtype=np.int64)
+
+        self.dummy_idx = count
+        self.embedding = keras.layers.Embedding(count + 1, embedding_dim)
+
+    def reset_parameters(self):
+        if self.embedding.built:
+            self.embedding.embeddings.assign(
+                keras.initializers.GlorotUniform()(self.embedding.embeddings.shape)
+            )
+
+    def __call__(self, node_type: str, batch: Optional[any] = None, **kwargs):
+        return self.call(node_type, batch)
+
+    def forward(self, node_type: str, batch: Optional[any] = None, **kwargs):
+        return self.call(node_type, batch)
+
+    def call(self, node_type: str, batch: Optional[any] = None):
+        r"""Returns the embeddings for the nodes in :obj:`batch` of type
+        :obj:`node_type`.
+        """
+        start = self.start[node_type]
+        end = self.end[node_type]
+        if batch is None:
+            batch = ops.arange(start, end, dtype="int64")
+        else:
+            batch = ops.cast(batch, "int64") + start
+        return self.embedding(batch)
+
+    def _pos_sample(self, batch):
+        batch_np = ops.convert_to_numpy(batch).astype(np.int64)
+        repeated = np.repeat(batch_np, self.walks_per_node)
+
+        all_walks = []
+        for node in repeated:
+            walk = [int(node)]
+            for i in range(self.walk_length):
+                edge_type = self.metapath[i % len(self.metapath)]
+                cur = walk[-1]
+                nbrs = (
+                    self.adj_dict[edge_type][cur]
+                    if cur < len(self.adj_dict[edge_type])
+                    else []
+                )
+                if len(nbrs) > 0:
+                    walk.append(nbrs[np.random.randint(len(nbrs))])
+                else:
+                    walk.append(self.dummy_idx)
+            all_walks.append(walk)
+
+        rw = np.array(all_walks, dtype=np.int64)
+        rw = rw + self.offset[None, :]
+        rw[rw > self.dummy_idx] = self.dummy_idx
+
+        walks = []
+        num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
+        for j in range(num_walks_per_rw):
+            walks.append(rw[:, j : j + self.context_size])
+        out = np.concatenate(walks, axis=0) if len(walks) > 0 else rw
+        return ops.convert_to_tensor(out, dtype="int64")
+
+    def _neg_sample(self, batch):
+        batch_np = ops.convert_to_numpy(batch).astype(np.int64)
+        repeated = np.repeat(
+            batch_np, self.walks_per_node * self.num_negative_samples
+        )
+
+        rws = [repeated]
+        for i in range(self.walk_length):
+            keys = self.metapath[i % len(self.metapath)]
+            num_nodes = self.num_nodes_dict[keys[-1]]
+            rand = np.random.randint(0, max(num_nodes, 1), size=len(repeated))
+            rws.append(rand)
+
+        rw = np.stack(rws, axis=-1)
+        rw = rw + self.offset[None, :]
+
+        walks = []
+        num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
+        for j in range(num_walks_per_rw):
+            walks.append(rw[:, j : j + self.context_size])
+        out = np.concatenate(walks, axis=0) if len(walks) > 0 else rw
+        return ops.convert_to_tensor(out, dtype="int64")
+
+    def loss(self, pos_rw, neg_rw):
+        r"""Computes the loss given positive and negative random walks."""
+        # Positive loss
+        start = pos_rw[:, 0]
+        rest = pos_rw[:, 1:]
+        pos_b = ops.shape(pos_rw)[0]
+
+        h_start = ops.reshape(self.embedding(start), (pos_b, 1, self.embedding_dim))
+        h_rest = ops.reshape(
+            self.embedding(ops.reshape(rest, (-1,))),
+            (pos_b, -1, self.embedding_dim),
+        )
+
+        out = ops.reshape(ops.sum(h_start * h_rest, axis=-1), (-1,))
+        pos_loss = -ops.mean(ops.log(ops.sigmoid(out) + self.EPS))
+
+        # Negative loss
+        start = neg_rw[:, 0]
+        rest = neg_rw[:, 1:]
+        neg_b = ops.shape(neg_rw)[0]
+
+        h_start = ops.reshape(self.embedding(start), (neg_b, 1, self.embedding_dim))
+        h_rest = ops.reshape(
+            self.embedding(ops.reshape(rest, (-1,))),
+            (neg_b, -1, self.embedding_dim),
+        )
+
+        out = ops.reshape(ops.sum(h_start * h_rest, axis=-1), (-1,))
+        neg_loss = -ops.mean(ops.log(1.0 - ops.sigmoid(out) + self.EPS))
+
+        return pos_loss + neg_loss

k3_node/models/mlp.py

@@ -0,0 +1,264 @@
+import re
+import inspect
+from typing import List, Optional, Union
+
+import keras
+from keras import ops
+
+import k3_node.layers.norm as norm_module
+
+
+def _normalize_string(s: str) -> str:
+    return re.sub(r"[_\-\s]", "", s).lower()
+
+
+def _normalization_resolver(query, *args, **kwargs):
+    if query is None:
+        return None
+    if not isinstance(query, str):
+        return query
+
+    norms = {
+        _normalize_string(name): cls
+        for name, cls in vars(norm_module).items()
+        if isinstance(cls, type)
+    }
+    key = _normalize_string(query)
+    if key in norms:
+        return norms[key](*args, **kwargs)
+    if key + "norm" in norms:
+        return norms[key + "norm"](*args, **kwargs)
+    if key.endswith("norm") and key[:-4] in norms:
+        return norms[key[:-4]](*args, **kwargs)
+    raise ValueError(f"Could not resolve normalization layer '{query}'")
+
+
+def _activation_resolver(act, **kwargs):
+    """Resolves Keras and PyG-style activation names ("relu", "leaky_relu", "LeakyReLU", ...)."""
+    if act is None:
+        return None
+    if isinstance(act, str):
+        name = act.lower().replace("_", "")
+        aliases = {"leakyrelu": "leaky_relu", "hardswish": "hard_swish", "hardsigmoid": "hard_sigmoid"}
+        act = keras.activations.get(aliases.get(name, act.lower()))
+    if kwargs:
+        import functools
+
+        return functools.partial(act, **kwargs)
+    return act
+
+
+class MLP(keras.Model):
+    r"""A Multi-Layer Perceptron (MLP) model.
+
+    There exists two ways to instantiate an `MLP`:
+
+    1. By specifying explicit channel sizes, e.g., `MLP([16, 32, 64, 128])`
+       creates a three-layer MLP with **differently** sized hidden layers.
+
+    2. By specifying fixed hidden channel sizes over a number of layers,
+       e.g., `MLP(in_channels=16, hidden_channels=32, out_channels=128,
+       num_layers=3)` creates a three-layer MLP with **equally** sized
+       hidden layers.
+
+    Args:
+        channel_list (List[int] or int, optional): List of input,
+            intermediate and output channels such that
+            `len(channel_list) - 1` denotes the number of layers of the
+            MLP. (default: `None`)
+        in_channels (int, optional): Size of each input sample. Will
+            override `channel_list`. (default: `None`)
+        hidden_channels (int, optional): Size of each hidden sample. Will
+            override `channel_list`. (default: `None`)
+        out_channels (int, optional): Size of each output sample. Will
+            override `channel_list`. (default: `None`)
+        num_layers (int, optional): The number of layers. Will override
+            `channel_list`. (default: `None`)
+        dropout (float or List[float], optional): Dropout probability of
+            each hidden embedding. (default: `0.`)
+        act (str or Callable, optional): The non-linear activation function
+            to use. (default: `"relu"`)
+        act_first (bool, optional): If set to `True`, activation is applied
+            before normalization. (default: `False`)
+        act_kwargs (dict, optional): Arguments passed to the activation function, e.g.
+            ``{"negative_slope": 0.2}`` for ``"leaky_relu"``. (default: `None`)
+        norm (str or Callable, optional): The normalization function to
+            use. (default: `"batch_norm"`)
+        norm_kwargs (dict, optional): Arguments passed to the respective
+            normalization function. (default: `None`)
+        plain_last (bool, optional): If set to `False`, will apply
+            non-linearity, normalization and dropout to the last layer as
+            well. (default: `True`)
+        bias (bool or List[bool], optional): If set to `False`, the module
+            will not learn additive biases. (default: `True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MLP
+
+        x = np.random.rand(10, 16).astype("float32")
+        mlp = MLP([16, 32, 32, 4])  # channel sizes: input, hidden, hidden, output
+        print(tuple(mlp(x).shape))  # (10, 4)
+
+        mlp = MLP(in_channels=16, hidden_channels=32, out_channels=4, num_layers=3, dropout=0.1)
+        print(tuple(mlp(x).shape))  # (10, 4)
+        ```
+    """
+    def __init__(
+        self,
+        channel_list: Optional[Union[List[int], int]] = None,
+        *args,
+        in_channels: Optional[int] = None,
+        hidden_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+        num_layers: Optional[int] = None,
+        dropout: Union[float, List[float]] = 0.0,
+        act="relu",
+        act_first: bool = False,
+        act_kwargs: Optional[dict] = None,
+        norm="batch_norm",
+        norm_kwargs: Optional[dict] = None,
+        plain_last: bool = True,
+        bias: Union[bool, List[bool]] = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if len(args) > 0:
+            if isinstance(channel_list, int):
+                channel_list = [channel_list] + list(args)
+            elif isinstance(channel_list, (list, tuple)):
+                channel_list = list(channel_list) + list(args)
+
+        if isinstance(channel_list, int):
+            in_channels = channel_list
+            channel_list = None
+
+        if in_channels is not None:
+            if num_layers is None:
+                raise ValueError("Argument `num_layers` must be given")
+            if num_layers > 1 and hidden_channels is None:
+                raise ValueError(
+                    f"Argument `hidden_channels` must be given for `num_layers={num_layers}`"
+                )
+            if out_channels is None:
+                raise ValueError("Argument `out_channels` must be given")
+
+            channel_list = [hidden_channels] * (num_layers - 1)
+            channel_list = [in_channels] + channel_list + [out_channels]
+
+        assert isinstance(channel_list, (tuple, list))
+        assert len(channel_list) >= 2
+        self.channel_list = list(channel_list)
+        self.in_channels = self.channel_list[0]
+        self.out_channels = self.channel_list[-1]
+
+        self.act = _activation_resolver(act, **(act_kwargs or {}))
+        self.act_first = act_first
+        self.plain_last = plain_last
+
+        if isinstance(dropout, float):
+            dropout = [dropout] * (len(channel_list) - 1)
+            if plain_last:
+                dropout[-1] = 0.0
+        if len(dropout) != len(channel_list) - 1:
+            raise ValueError(
+                f"Number of dropout values provided ({len(dropout)}) does not "
+                f"match the number of layers specified ({len(channel_list) - 1})"
+            )
+        self.dropout_rate = dropout
+
+        if isinstance(bias, bool):
+            bias = [bias] * (len(channel_list) - 1)
+        if len(bias) != len(channel_list) - 1:
+            raise ValueError(
+                f"Number of bias values provided ({len(bias)}) does not match "
+                f"the number of layers specified ({len(channel_list) - 1})"
+            )
+
+        self.lins = []
+        for in_c, out_c, _bias in zip(channel_list[:-1], channel_list[1:], bias):
+            lin = keras.layers.Dense(out_c, use_bias=_bias)
+            lin.build((None, in_c))
+            self.lins.append(lin)
+
+        self.norms = []
+        iterator = channel_list[1:-1] if plain_last else channel_list[1:]
+        for hc in iterator:
+            norm_layer = _normalization_resolver(norm, hc, **(norm_kwargs or {}))
+            if norm_layer is not None and hasattr(norm_layer, "build") and not norm_layer.built:
+                norm_layer.build((None, hc))
+            self.norms.append(norm_layer)
+
+        self.dropouts = [keras.layers.Dropout(p) if p > 0.0 else None for p in self.dropout_rate]
+
+        self.supports_norm_batch = False
+        if len(self.norms) > 0 and self.norms[0] is not None:
+            norm_params = inspect.signature(self.norms[0].call).parameters
+            self.supports_norm_batch = "batch" in norm_params
+        # Norms such as BatchNorm treat ``training=None`` as training mode, so pass it explicitly.
+        self.supports_norm_training = False
+        if len(self.norms) > 0 and self.norms[0] is not None:
+            self.supports_norm_training = "training" in inspect.signature(self.norms[0].call).parameters
+
+    @property
+    def num_layers(self) -> int:
+        r"""The number of layers."""
+        return len(self.channel_list) - 1
+
+    def build(self, input_shape=None):
+        for lin in self.lins:
+            if hasattr(lin, "built") and not lin.built:
+                lin.build(input_shape)
+        norm_iter = self.channel_list[1:-1] if self.plain_last else self.channel_list[1:]
+        for norm, hc in zip(self.norms, norm_iter):
+            if norm is not None and hasattr(norm, "built") and not norm.built:
+                norm.build((None, hc))
+        for drop in self.dropouts:
+            if drop is not None and hasattr(drop, "built") and not drop.built:
+                drop.build(input_shape)
+        self.built = True
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for lin in self.lins:
+            if hasattr(lin, "kernel_initializer") and lin.kernel is not None:
+                lin.kernel.assign(lin.kernel_initializer(ops.shape(lin.kernel)))
+                if lin.bias is not None:
+                    lin.bias.assign(lin.bias_initializer(ops.shape(lin.bias)))
+        for norm in self.norms:
+            if hasattr(norm, "reset_parameters"):
+                norm.reset_parameters()
+
+    def call(self, x, batch=None, batch_size=None, return_emb=None, training=None):
+        emb = None
+
+        # If `plain_last=True`, `len(norms) == len(lins) - 1`, thus skipping
+        # execution of the last layer inside the loop.
+        for i, (lin, norm) in enumerate(zip(self.lins, self.norms)):
+            x = lin(x)
+            if self.act is not None and self.act_first:
+                x = self.act(x)
+            if norm is not None:
+                norm_kwargs = {"training": training} if self.supports_norm_training else {}
+                if self.supports_norm_batch:
+                    x = norm(x, batch, batch_size, **norm_kwargs)
+                else:
+                    x = norm(x, **norm_kwargs)
+            if self.act is not None and not self.act_first:
+                x = self.act(x)
+            if self.dropouts[i] is not None:
+                x = self.dropouts[i](x, training=training)
+            if isinstance(return_emb, bool) and return_emb is True:
+                emb = x
+
+        if self.plain_last:
+            x = self.lins[-1](x)
+            if self.dropouts[-1] is not None:
+                x = self.dropouts[-1](x, training=training)
+
+        return (x, emb) if isinstance(return_emb, bool) else x
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({str(self.channel_list)[1:-1]})"