Geometric EL Embeddings: Training Animation

All three geometric EL models embed ontology classes as 2-D shapes when trained with embed_dim=2, letting us watch the geometry evolve directly — no dimensionality reduction needed.

  • ELEm — classes as circles (centre class_embed, radius class_rad)

  • ELBE — classes as axis-aligned rectangles (centre class_embed, half-extents class_offset)

  • Box²EL — classes as axis-aligned rectangles (centre class_center, half-extents class_offset); roles as pairs of boxes — head box (solid, soft fill) and tail box (dashed, soft fill)

This example trains all three models on the Family ontology and produces an interactive animation showing how the shapes evolve across epochs.

Imports and JVM initialisation

import mowl
mowl.init_jvm("10g")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.animation as animation
import numpy as np
import torch as th
from IPython.display import HTML

from mowl.datasets.builtin import FamilyDataset
from mowl.models import ELEmbeddings, ELBE, BoxSquaredEL
from mowl.visualization.el.base import _local_name

Dataset and entity selection

dataset = FamilyDataset()

Training parameters

EPOCHS = 100
SNAPSHOT_EVERY = 5
N_CLASSES = 6

Snapshot helpers

We store raw numpy arrays per snapshot and recreate patches at animation time.

def _indices(model, n):
    items = list(model.class_index_dict.items())[:n]
    iris = [iri for iri, _ in items]
    idxs = th.tensor([idx for _, idx in items])
    return iris, idxs


def _role_indices(model):
    items = list(model.object_property_index_dict.items())
    iris = [iri for iri, _ in items]
    idxs = th.tensor([idx for _, idx in items])
    return iris, idxs


def _elem_snapshot(model, idxs):
    with th.no_grad():
        centers = model.module.class_embed.weight[idxs].cpu().numpy()
        radii   = np.abs(model.module.class_rad.weight[idxs].cpu().numpy())
    return centers, radii


def _box_snapshot_elbe(model, idxs):
    with th.no_grad():
        centers = model.module.class_embed.weight[idxs].cpu().numpy()
        halves  = np.abs(model.module.class_offset.weight[idxs].cpu().numpy())
    return centers, halves


def _box_snapshot_box2(model, class_idxs, role_idxs):
    with th.no_grad():
        centers    = model.module.class_center.weight[class_idxs].cpu().numpy()
        halves     = np.abs(model.module.class_offset.weight[class_idxs].cpu().numpy())
        head_cs    = model.module.head_center.weight[role_idxs].cpu().numpy()
        head_hs    = np.abs(model.module.head_offset.weight[role_idxs].cpu().numpy())
        tail_cs    = model.module.tail_center.weight[role_idxs].cpu().numpy()
        tail_hs    = np.abs(model.module.tail_offset.weight[role_idxs].cpu().numpy())
    return centers, halves, head_cs, head_hs, tail_cs, tail_hs


def make_callback(snapshots, getter, *getter_args):
    def callback(epoch, model):
        if epoch % SNAPSHOT_EVERY == 0:
            snapshots.append(getter(model, *getter_args))
    return callback

Train ELEm

elem_model = ELEmbeddings(dataset, embed_dim=2, learning_rate=0.01, margin=-0.1)
iris, idxs = _indices(elem_model, N_CLASSES)
labels = [_local_name(iri) for iri in iris]

elem_snapshots = []
elem_model.train(epochs=EPOCHS,
                 epoch_callback=make_callback(elem_snapshots, _elem_snapshot, idxs))

Train ELBE

elbe_model = ELBE(dataset, embed_dim=2, learning_rate=0.01, margin=-0.1)
elbe_snapshots = []
elbe_model.train(epochs=EPOCHS,
                 epoch_callback=make_callback(elbe_snapshots, _box_snapshot_elbe, idxs))

Train Box²EL

box2_model = BoxSquaredEL(dataset, embed_dim=2, learning_rate=0.01, margin=-0.1)
role_iris, role_idxs = _role_indices(box2_model)
role_labels = [_local_name(iri) for iri in role_iris]

box2_snapshots = []
box2_model.train(epochs=EPOCHS,
                 epoch_callback=make_callback(box2_snapshots, _box_snapshot_box2, idxs, role_idxs))

Build and display the animation

colors = plt.cm.tab10(np.linspace(0, 0.6, N_CLASSES))
role_colors = plt.cm.Set1(np.linspace(0, 0.5, max(1, len(role_iris))))


def _axis_limits(all_snaps, pad=0.3):
    """Compute unified axis limits across all models and all snapshots."""
    xs, ys = [], []
    for snaps in all_snaps:
        for snap in snaps:
            centers, sizes = snap[0], snap[1]
            if sizes.shape[1] == 1:
                r = sizes[:, 0]
                xs.extend((centers[:, 0] - r).tolist())
                xs.extend((centers[:, 0] + r).tolist())
                ys.extend((centers[:, 1] - r).tolist())
                ys.extend((centers[:, 1] + r).tolist())
            else:
                xs.extend((centers[:, 0] - sizes[:, 0]).tolist())
                xs.extend((centers[:, 0] + sizes[:, 0]).tolist())
                ys.extend((centers[:, 1] - sizes[:, 1]).tolist())
                ys.extend((centers[:, 1] + sizes[:, 1]).tolist())
            # Box²EL snaps carry role box extents in positions 2-5
            if len(snap) == 6:
                head_cs, head_hs, tail_cs, tail_hs = snap[2], snap[3], snap[4], snap[5]
                for cs, hs in [(head_cs, head_hs), (tail_cs, tail_hs)]:
                    xs.extend((cs[:, 0] - hs[:, 0]).tolist())
                    xs.extend((cs[:, 0] + hs[:, 0]).tolist())
                    ys.extend((cs[:, 1] - hs[:, 1]).tolist())
                    ys.extend((cs[:, 1] + hs[:, 1]).tolist())
    return min(xs) - pad, max(xs) + pad, min(ys) - pad, max(ys) + pad


xmin, xmax, ymin, ymax = _axis_limits([elem_snapshots, elbe_snapshots, box2_snapshots])

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
model_titles = ["ELEm (circles)", "ELBE (boxes)", "Box²EL (boxes + role pairs)"]
all_snaps = [elem_snapshots, elbe_snapshots, box2_snapshots]


def draw_frame(frame_idx):
    for ax, snaps, title in zip(axes, all_snaps, model_titles):
        ax.cla()
        ax.set_title(title, fontsize=11)
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_aspect("equal")
        ax.grid(True, alpha=0.3)

        snap = snaps[frame_idx]
        centers, sizes = snap[0], snap[1]

        # Draw class shapes
        for i, (cx, cy) in enumerate(centers):
            color = colors[i]
            if sizes.shape[1] == 1:
                r = float(sizes[i, 0])
                patch = mpatches.Circle((cx, cy), r, fill=False,
                                        edgecolor=color, linewidth=1.5)
            else:
                hx, hy = float(sizes[i, 0]), float(sizes[i, 1])
                patch = mpatches.Rectangle(
                    (cx - hx, cy - hy), 2 * hx, 2 * hy,
                    fill=False, edgecolor=color, linewidth=1.5)
            ax.add_patch(patch)
            ax.annotate(labels[i], (cx, cy), ha="center", va="center",
                        fontsize=7, color=color)

        # Draw role boxes for Box²EL: head=solid soft fill, tail=dashed soft fill
        if len(snap) == 6:
            head_cs, head_hs, tail_cs, tail_hs = snap[2], snap[3], snap[4], snap[5]
            for j, rl in enumerate(role_labels):
                rc = role_colors[j]

                hcx, hcy = float(head_cs[j, 0]), float(head_cs[j, 1])
                hhx, hhy = float(head_hs[j, 0]), float(head_hs[j, 1])
                ax.add_patch(mpatches.Rectangle(
                    (hcx - hhx, hcy - hhy), 2 * hhx, 2 * hhy,
                    fill=True, facecolor=rc, alpha=0.15,
                    edgecolor=rc, linestyle="solid", linewidth=1.5))
                ax.annotate(f"{rl}(head)", (hcx, hcy), ha="center", va="center",
                            fontsize=6, color=rc)

                tcx, tcy = float(tail_cs[j, 0]), float(tail_cs[j, 1])
                thx, thy = float(tail_hs[j, 0]), float(tail_hs[j, 1])
                ax.add_patch(mpatches.Rectangle(
                    (tcx - thx, tcy - thy), 2 * thx, 2 * thy,
                    fill=True, facecolor=rc, alpha=0.15,
                    edgecolor=rc, linestyle="dashed", linewidth=1.5))
                ax.annotate(f"{rl}(tail)", (tcx, tcy), ha="center", va="center",
                            fontsize=6, color=rc)

    fig.suptitle(f"Epoch {frame_idx * SNAPSHOT_EVERY}", fontsize=13)


n_frames = len(elem_snapshots)
ani = animation.FuncAnimation(fig, draw_frame, frames=n_frames, interval=150)

Save as an animated GIF that autoplays in browsers and docs pages. When run as a standalone script, the GIF is saved next to the docs images. In a Sphinx Gallery build the file is pre-committed; __file__ is not available.

try:
    import os as _os
    _gif_path = _os.path.join(
        _os.path.dirname(_os.path.abspath(__file__)),
        "../../docs/source/examples/elmodels/images/el_geometry_training.gif"
    )
    ani.save(_gif_path, writer="pillow", fps=6)
except NameError:
    pass
ELEm, ELBE, and Box²EL shapes optimizing on the Family ontology across epochs

Estimated memory usage: 0 MB

Gallery generated by Sphinx-Gallery