Skip to content
All library documents

Building Macro Group Adjacency Masks for Asset Attention

Code Machine Learning for Trading

Summary

This utility constructs a Boolean asset-to-asset adjacency matrix from asset group labels and optional links between groups. Assets assigned to the same macro group are connected to one another, and each supplied cross-group edge connects all members of both groups in both directions. Optional self-links allow each asset to attend to itself. Assets without a mapping are assigned to an unknown group, so they are connected to other unmapped assets.

A second function converts the adjacency matrix into an attention mask by marking non-adjacent pairs as disallowed. The code checks that the input matrix is square before conversion. This provides a simple way to encode a macro grouping prior into an attention model, but it does not estimate groups or relationships from market data, and its usefulness depends on the quality and currency of the supplied mapping and cross-group edges. The document gives implementation details rather than empirical results or a trading evaluation.

Key ideas

  • Assets in the same group receive mutual connections in the adjacency matrix.
  • Cross-group links connect every asset in one group to every asset in the other, symmetrically.
  • Self-connections can be included or omitted through a parameter.
  • Unmapped assets share an unknown group and therefore connect to one another.
  • The attention mask marks every non-adjacent asset pair as disallowed.

Tags

Full text
# graph.py


```py
"""Macroeconomic graph prior utilities.

Builds an asset-level adjacency matrix from macro group labels and
cross-group edges for adjacency-masked attention in DeePM.
"""

from __future__ import annotations

from collections.abc import Iterable, Sequence
from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True, slots=True)
class MacroGraph:
    """Asset-level macro graph."""

    assets: list[str]
    groups: list[str]
    adjacency: np.ndarray  # Boolean (N, N)


def build_macro_adjacency(
    *,
    assets: Sequence[str],
    asset_to_group: dict[str, str],
    cross_group_edges: Iterable[tuple[str, str]] = (),
    include_self_loops: bool = True,
) -> MacroGraph:
    """Build a boolean adjacency matrix from group labels and group edges.

    Parameters
    ----------
    assets: Asset identifiers matching the price panel columns.
    asset_to_group: Mapping from asset -> macro group label.
    cross_group_edges: Undirected (group_a, group_b) edges.
    include_self_loops: If True, sets A[i,i] = True.
    """
    assets_list = [str(a) for a in assets]
    groups = [str(asset_to_group.get(a, "UNKNOWN")) for a in assets_list]

    n = len(assets_list)
    adj = np.zeros((n, n), dtype=bool)

    if include_self_loops:
        np.fill_diagonal(adj, True)

    group_to_indices: dict[str, list[int]] = {}
    for i, g in enumerate(groups):
        group_to_indices.setdefault(g, []).append(i)

    for indices in group_to_indices.values():
        idx = np.array(indices, dtype=int)
        adj[np.ix_(idx, idx)] = True

    for g1, g2 in cross_group_edges:
        idx1 = group_to_indices.get(g1, [])
        idx2 = group_to_indices.get(g2, [])
        if not idx1 or not idx2:
            continue
        a = np.array(idx1, dtype=int)
        b = np.array(idx2, dtype=int)
        adj[np.ix_(a, b)] = True
        adj[np.ix_(b, a)] = True

    return MacroGraph(assets=assets_list, groups=groups, adjacency=adj)


def adjacency_to_attn_mask(adjacency: np.ndarray) -> np.ndarray:
    """Convert boolean adjacency to attention mask (True = disallowed)."""
    if adjacency.ndim != 2 or adjacency.shape[0] != adjacency.shape[1]:
        raise ValueError("adjacency must be square (N,N)")
    return ~adjacency

```

Shown in full with attribution under the source's licence. Licence: MIT

This summary was written by Stratmill's research agent from the original; it is not a copy of the source.