Does an LLM Store Cities on a Globe?¶

When a language model reads "I live in Tokyo", does the hidden state at the city token encode where Tokyo is on Earth — and with what geometry?

The pipeline: label "I live in {city}" sentences with (lat, lon), capture the residual stream at the city token, then fit three competing SMDS geometries — spherical, cylindrical and flat — and score them on cities they did not see. That comparison is the point: SMDS does not ask "is there structure?" but "which of these shapes explains it best?"

Dependencies. nnsight and torch are not part of the base smds install: uv add nnsight torch (or pip install nnsight torch).

In [25]:

Copied!





from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.io as pio
import torch
from nnsight import LanguageModel
from numpy.typing import NDArray
from scipy.linalg import orthogonal_procrustes
from sklearn.model_selection import StratifiedKFold, cross_validate, train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm

from smds import SupervisedMDS

# mkdocs-jupyter exports the docs with nbconvert, which has no renderer for plotly's own mime
# bundle and silently drops the figure. This renderer emits self-contained HTML with a plotly.js
# CDN tag instead, so `fig.show()` output survives the export and renders on the docs site.
pio.renderers.default = "notebook_connected"
from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.io as pio
import torch
from nnsight import LanguageModel
from numpy.typing import NDArray
from scipy.linalg import orthogonal_procrustes
from sklearn.model_selection import StratifiedKFold, cross_validate, train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm

from smds import SupervisedMDS

# mkdocs-jupyter exports the docs with nbconvert, which has no renderer for plotly's own mime
# bundle and silently drops the figure. This renderer emits self-contained HTML with a plotly.js
# CDN tag instead, so `fig.show()` output survives the export and renders on the docs site.
pio.renderers.default = "notebook_connected"

Configuration¶

In [26]:

Copied!





# Gated repo -- `huggingface-cli login` (or HF_TOKEN) first. Any causal LM works.
MODEL_ID = "meta-llama/Llama-3.2-1B"
BATCH_SIZE = 128
SEED = 123
PROMPT_TEMPLATE = "I live in {city}"

# Derived table shipped with the notebook; see the provenance note in Stage 1.
CITIES_CSV = "world_cities_continents.csv"
# The one knob on the city list: dataset size against how likely the model is to know a place.
# 50k leaves 5,705 cities.
MIN_POPULATION = 50_000
# Prototyping knob: cap the dataset at roughly this many cities (None = use all of them).
MAX_CITIES = None

TEST_SIZE = 0.5  # the split behind the 3D figure: half trains the projection, half is drawn held out
N_FOLDS = 20  # the stress table cross-validates instead of trusting one split
N_COMPONENTS = 3  # same dimensional budget for every geometry -- what makes their stress comparable
ALPHA = 1.0  # ridge on W; non-zero matters, d >> N would otherwise interpolate the training half

# Memory, not cores, is what caps the cross-validation. Every fold holds several n_train x n_train
# matrices -- the classical-MDS eigendecomposition peaks at ~2.6 GB for 5.7k cities, ~3.8 GB once the
# train fold is scored -- and each parallel worker pays that in full. Budget ~4 GB per worker against
# total RAM: `n_jobs=-1` asks 12 workers for ~45 GB, the OS kills them, and joblib reports
# "A worker stopped while some jobs were given to the executor".
N_JOBS = 3

DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")

# Cache key includes MAX_CITIES: a subsampled run must not reuse the full run's activations.
CACHE_DIR = Path("activations")
CACHE_TAG = MODEL_ID.replace("/", "_") + ("" if MAX_CITIES is None else f"_n{MAX_CITIES}")
CACHE_PATH = CACHE_DIR / f"city_manifold_{CACHE_TAG}.npz"

print(f"Running on {DEVICE}")
# Gated repo -- `huggingface-cli login` (or HF_TOKEN) first. Any causal LM works.
MODEL_ID = "meta-llama/Llama-3.2-1B"
BATCH_SIZE = 128
SEED = 123
PROMPT_TEMPLATE = "I live in {city}"

# Derived table shipped with the notebook; see the provenance note in Stage 1.
CITIES_CSV = "world_cities_continents.csv"
# The one knob on the city list: dataset size against how likely the model is to know a place.
# 50k leaves 5,705 cities.
MIN_POPULATION = 50_000
# Prototyping knob: cap the dataset at roughly this many cities (None = use all of them).
MAX_CITIES = None

TEST_SIZE = 0.5  # the split behind the 3D figure: half trains the projection, half is drawn held out
N_FOLDS = 20  # the stress table cross-validates instead of trusting one split
N_COMPONENTS = 3  # same dimensional budget for every geometry -- what makes their stress comparable
ALPHA = 1.0  # ridge on W; non-zero matters, d >> N would otherwise interpolate the training half

# Memory, not cores, is what caps the cross-validation. Every fold holds several n_train x n_train
# matrices -- the classical-MDS eigendecomposition peaks at ~2.6 GB for 5.7k cities, ~3.8 GB once the
# train fold is scored -- and each parallel worker pays that in full. Budget ~4 GB per worker against
# total RAM: `n_jobs=-1` asks 12 workers for ~45 GB, the OS kills them, and joblib reports
# "A worker stopped while some jobs were given to the executor".
N_JOBS = 3

DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")

# Cache key includes MAX_CITIES: a subsampled run must not reuse the full run's activations.
CACHE_DIR = Path("activations")
CACHE_TAG = MODEL_ID.replace("/", "_") + ("" if MAX_CITIES is None else f"_n{MAX_CITIES}")
CACHE_PATH = CACHE_DIR / f"city_manifold_{CACHE_TAG}.npz"

print(f"Running on {DEVICE}")

Running on mps

Stage 1 — Dataset¶

One row per sentence, labelled with the city's latitude, longitude and continent. The sentence ends on the city name, so the capture anchor is always the last token.

Data provenance. world_cities_continents.csv is a derivative of the simplemaps World Cities Database, Basic edition (CC BY 4.0), post-processed to attach a continent to every city. The post-processing kept administrative centres plus US/Canadian cities above 100k, joined a country→continent lookup, filled the ~444 cities the join missed by majority vote over their 3 nearest labelled neighbours, and collapsed duplicate city names to the most populous entry.

That leaves 14,776 cities; MIN_POPULATION thins them to 5,705 spanning the full ±180° of longitude. Set MAX_CITIES to thin further — the cut is stratified by continent — for a quicker run.

In [27]:

Copied!





# A derivative of the simplemaps World Cities Database (CC BY 4.0), already filtered to
# administrative centres and annotated with continents -- see the provenance note above.
df = pd.read_csv(CITIES_CSV)
df = df[df["population"] >= MIN_POPULATION]

df = df[["city", "country", "continent", "lat", "lon"]].sort_values("city").reset_index(drop=True)
df["sentence"] = df["city"].map(lambda city: PROMPT_TEMPLATE.format(city=city))
df = df[["sentence", "city", "country", "lat", "lon", "continent"]]

if MAX_CITIES is not None:
    # Stratified thinning: keep each continent's share, but never fewer than N_FOLDS of it.
    frac = MAX_CITIES / len(df)
    thinned = [
        group.sample(min(len(group), max(N_FOLDS, round(frac * len(group)))), random_state=SEED)
        for _, group in df.groupby("continent")
    ]
    df = pd.concat(thinned).sort_values("city").reset_index(drop=True)

print(f"{len(df)} cities across {df['continent'].nunique()} continents")
print(f"latitude {df['lat'].min():.0f}..{df['lat'].max():.0f}   longitude {df['lon'].min():.0f}..{df['lon'].max():.0f}")
print(df["continent"].value_counts().to_string())
df.head()
# A derivative of the simplemaps World Cities Database (CC BY 4.0), already filtered to
# administrative centres and annotated with continents -- see the provenance note above.
df = pd.read_csv(CITIES_CSV)
df = df[df["population"] >= MIN_POPULATION]

df = df[["city", "country", "continent", "lat", "lon"]].sort_values("city").reset_index(drop=True)
df["sentence"] = df["city"].map(lambda city: PROMPT_TEMPLATE.format(city=city))
df = df[["sentence", "city", "country", "lat", "lon", "continent"]]

if MAX_CITIES is not None:
    # Stratified thinning: keep each continent's share, but never fewer than N_FOLDS of it.
    frac = MAX_CITIES / len(df)
    thinned = [
        group.sample(min(len(group), max(N_FOLDS, round(frac * len(group)))), random_state=SEED)
        for _, group in df.groupby("continent")
    ]
    df = pd.concat(thinned).sort_values("city").reset_index(drop=True)

print(f"{len(df)} cities across {df['continent'].nunique()} continents")
print(f"latitude {df['lat'].min():.0f}..{df['lat'].max():.0f}   longitude {df['lon'].min():.0f}..{df['lon'].max():.0f}")
print(df["continent"].value_counts().to_string())
df.head()

5705 cities across 6 continents
latitude -55..69   longitude -158..178
continent
Asia             1843
Europe           1267
South America    1001
North America     972
Africa            600
Oceania            22

Out[27]:

	sentence	city	country	lat	lon	continent
0	I live in A Coruña	A Coruña	Spain	43.3650	-8.4100	Europe
1	I live in A Yun Pa	A Yun Pa	Vietnam	13.3939	108.4408	Asia
2	I live in Aachen	Aachen	Germany	50.7756	6.0836	Europe
3	I live in Aalborg	Aalborg	Denmark	57.0337	9.9166	Europe
4	I live in Aalen	Aalen	Germany	48.8333	10.1000	Europe

Stage 2 — Activation capture¶

Inside model.trace, nnsight saves the residual-stream output of the middle layer without altering the forward pass. Batching needs padding, so position -1 is not the city token for every row and each row's anchor is read from its own last non-pad index instead.

That produces the two inputs SMDS needs — $X$, one activation per sentence, and $y$, the matching (latitude, longitude) pairs — cached to .npz so the model is loaded only once.

In [28]:

Copied!





def resolve_residual_layers(model: LanguageModel):
    """Return the list of residual-stream blocks for GPT-style or LLaMA-style checkpoints."""
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    raise ValueError("Unsupported model structure for residual stream access.")


def capture_city_token_activations(sentences: list[str]) -> NDArray[np.float32]:
    """Trace the sentences in batches, returning the hidden state at each city token."""
    model = LanguageModel(MODEL_ID, device_map=DEVICE, dtype=torch.float32, dispatch=True)
    tokenizer = model.tokenizer
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"  # trailing pads are invisible to earlier positions under causal attention

    layers = resolve_residual_layers(model)
    mid_layer = len(layers) // 2
    print(f"Capturing layer {mid_layer} of {len(layers)} (exact middle of the network)")

    activations: list[NDArray[np.float32]] = []

    for start in tqdm(range(0, len(sentences), BATCH_SIZE)):
        batch = tokenizer(sentences[start : start + BATCH_SIZE], return_tensors="pt", padding=True).to(DEVICE)

        with model.trace(batch):
            hidden = layers[mid_layer].output.save()

        # Older transformers versions return (hidden_states, ...) from a decoder layer.
        hidden = hidden[0] if isinstance(hidden, tuple) else hidden
        hidden = hidden.detach().float().cpu()  # (B, T, d)

        # Per row, the last non-pad position: the final sub-token of the city name.
        last_index = batch["attention_mask"].sum(dim=1).cpu() - 1
        rows = torch.arange(hidden.shape[0])

        activations.append(hidden[rows, last_index].numpy())

    return np.concatenate(activations).astype(np.float32)


if CACHE_PATH.exists():
    X = np.load(CACHE_PATH)["X"]
    print(f"Loaded cached activations from {CACHE_PATH}")
else:
    X = capture_city_token_activations(df["sentence"].tolist())
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    np.savez(CACHE_PATH, X=X)
    print(f"Saved activations to {CACHE_PATH}")

y = df[["lat", "lon"]].to_numpy(dtype=np.float64)

print(f"X: {X.shape}   (N x d)")
print(f"y: {y.shape}   (N x 2, degrees)")
def resolve_residual_layers(model: LanguageModel):
    """Return the list of residual-stream blocks for GPT-style or LLaMA-style checkpoints."""
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    raise ValueError("Unsupported model structure for residual stream access.")


def capture_city_token_activations(sentences: list[str]) -> NDArray[np.float32]:
    """Trace the sentences in batches, returning the hidden state at each city token."""
    model = LanguageModel(MODEL_ID, device_map=DEVICE, dtype=torch.float32, dispatch=True)
    tokenizer = model.tokenizer
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"  # trailing pads are invisible to earlier positions under causal attention

    layers = resolve_residual_layers(model)
    mid_layer = len(layers) // 2
    print(f"Capturing layer {mid_layer} of {len(layers)} (exact middle of the network)")

    activations: list[NDArray[np.float32]] = []

    for start in tqdm(range(0, len(sentences), BATCH_SIZE)):
        batch = tokenizer(sentences[start : start + BATCH_SIZE], return_tensors="pt", padding=True).to(DEVICE)

        with model.trace(batch):
            hidden = layers[mid_layer].output.save()

        # Older transformers versions return (hidden_states, ...) from a decoder layer.
        hidden = hidden[0] if isinstance(hidden, tuple) else hidden
        hidden = hidden.detach().float().cpu()  # (B, T, d)

        # Per row, the last non-pad position: the final sub-token of the city name.
        last_index = batch["attention_mask"].sum(dim=1).cpu() - 1
        rows = torch.arange(hidden.shape[0])

        activations.append(hidden[rows, last_index].numpy())

    return np.concatenate(activations).astype(np.float32)


if CACHE_PATH.exists():
    X = np.load(CACHE_PATH)["X"]
    print(f"Loaded cached activations from {CACHE_PATH}")
else:
    X = capture_city_token_activations(df["sentence"].tolist())
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    np.savez(CACHE_PATH, X=X)
    print(f"Saved activations to {CACHE_PATH}")

y = df[["lat", "lon"]].to_numpy(dtype=np.float64)

print(f"X: {X.shape}   (N x d)")
print(f"y: {y.shape}   (N x 2, degrees)")

Loaded cached activations from activations/city_manifold_meta-llama_Llama-3.2-1B.npz
X: (5705, 2048)   (N x d)
y: (5705, 2)   (N x 2, degrees)

Stage 3 — Fitting three competing geometries¶

Half the cities train the projection, half are held out. SMDS learns a linear map with no intercept, so the features are standardised first, fitted on the training half only.

Geometry	Ideal distance between two cities	What it assumes
Spherical	chord length on the unit sphere	a globe: both coordinates wrap
Cylindrical	Euclidean on a cylinder	longitude wraps, latitude does not
Flat (linear)	Euclidean in the (lat, lon) plane	a flat map: no wrap at the antimeridian

All three manifolds are computed from the same (lat, lon) labels, and all three get n_components=3.

In [29]:

Copied!





idx_train, idx_test = train_test_split(
    np.arange(len(df)),
    test_size=TEST_SIZE,
    random_state=SEED,
    stratify=df["continent"],
)

scaler = StandardScaler().fit(X[idx_train])
X_scaled = scaler.transform(X)

X_train, y_train = X_scaled[idx_train], y[idx_train]
X_test, y_test = X_scaled[idx_test], y[idx_test]

models: dict[str, SupervisedMDS] = {
    "spherical": SupervisedMDS(manifold="spherical", n_components=N_COMPONENTS, alpha=ALPHA),
    "flat (linear)": SupervisedMDS(manifold="euclidean", n_components=N_COMPONENTS, alpha=ALPHA),
    "cylindrical": SupervisedMDS(manifold="cylindrical", n_components=N_COMPONENTS, alpha=ALPHA),
}

for name, model in models.items():
    model.fit(X_train, y_train)
    n_comp = model.parametrization_fitted_.n_components
    print(f"{name:<14} fitted   n_components={n_comp}   W: {model.W_.shape}")

print(f"\nTrain: {len(idx_train)} cities   Test: {len(idx_test)} cities")
idx_train, idx_test = train_test_split(
    np.arange(len(df)),
    test_size=TEST_SIZE,
    random_state=SEED,
    stratify=df["continent"],
)

scaler = StandardScaler().fit(X[idx_train])
X_scaled = scaler.transform(X)

X_train, y_train = X_scaled[idx_train], y[idx_train]
X_test, y_test = X_scaled[idx_test], y[idx_test]

models: dict[str, SupervisedMDS] = {
    "spherical": SupervisedMDS(manifold="spherical", n_components=N_COMPONENTS, alpha=ALPHA),
    "flat (linear)": SupervisedMDS(manifold="euclidean", n_components=N_COMPONENTS, alpha=ALPHA),
    "cylindrical": SupervisedMDS(manifold="cylindrical", n_components=N_COMPONENTS, alpha=ALPHA),
}

for name, model in models.items():
    model.fit(X_train, y_train)
    n_comp = model.parametrization_fitted_.n_components
    print(f"{name:<14} fitted   n_components={n_comp}   W: {model.W_.shape}")

print(f"\nTrain: {len(idx_train)} cities   Test: {len(idx_test)} cities")

spherical      fitted   n_components=3   W: (3, 2048)
flat (linear)  fitted   n_components=3   W: (3, 2048)
cylindrical    fitted   n_components=3   W: (3, 2048)

Train: 2852 cities   Test: 2853 cities

Stage 4 — Visualising the embedding¶

Every city, train and held out alike, is projected through the spherical model's learned $W$ into three dimensions. If the hypothesis holds, the points sit near the surface of a unit sphere with continents grouped as they are on Earth. Colour encodes continent, open diamonds mark held-out cities, hovering names the city, and Lisbon is called out as a landmark.

Note: An SMDS embedding is only pinned down up to an isometry, so two linear corrections are applied to recover the correct centring and orientation of the globe.

In [30]:

Copied!





# Fixed categorical order -- hues are assigned per continent and never cycled.
CONTINENT_COLORS = {
    "Europe": "#2a78d6",
    "Africa": "#eb6834",
    "Asia": "#1baf7a",
    "North America": "#eda100",
    "South America": "#e87ba4",
    "Oceania": "#4a3aa7",
}

def fit_sphere(points: NDArray[np.float64]) -> tuple[NDArray[np.float64], float]:
    """Least-squares sphere through the points, as (centre, radius)."""
    design = np.c_[2 * points, np.ones(len(points))]
    solution, *_ = np.linalg.lstsq(design, (points**2).sum(axis=1), rcond=None)
    centre = solution[:3]
    radius = float(np.sqrt(solution[3] + centre @ centre))
    return centre, radius


def globe_coords(coordinates: NDArray[np.float64]) -> NDArray[np.float64]:
    """Map (lat, lon) in degrees onto the unit sphere, z pointing at the North Pole."""
    lat, lon = np.radians(coordinates[:, 0]), np.radians(coordinates[:, 1])
    return np.stack([np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], axis=1)


spherical = models["spherical"]

# The template's own sphere. Classical MDS centres the embedding on the centroid of the sampled
# cities, which is not the centre of the globe, so the fitted centre is offset from the origin.
centre, radius = fit_sphere(spherical.Y_)
proj = spherical.transform(X_scaled) - centre

# An SMDS embedding is only defined up to an isometry: rotations and reflections leave every pairwise
# distance untouched, so the raw axes point in arbitrary directions and the globe comes out mirrored.
# Procrustes onto the true coordinates of the *training* cities puts the cloud in the Earth frame --
# z through the North Pole, x through (0 N, 0 E) -- which is what makes the picture readable.
rotation, _ = orthogonal_procrustes(proj[idx_train], globe_coords(y[idx_train]))
proj = proj @ rotation

print(f"template sphere: fitted centre {np.round(centre, 3)}, radius {radius:.3f}")
print(f"alignment to the Earth frame: {'reflection + rotation' if np.linalg.det(rotation) < 0 else 'rotation'}")

plot_df = df.assign(
    x=proj[:, 0],
    y_=proj[:, 1],
    z=proj[:, 2],
    split=np.where(np.isin(np.arange(len(df)), idx_train), "train", "held-out"),
)

u, v = np.mgrid[0 : 2 * np.pi : 60j, 0 : np.pi : 30j]
fig = go.Figure(
    go.Surface(
        x=radius * np.cos(u) * np.sin(v),
        y=radius * np.sin(u) * np.sin(v),
        z=radius * np.cos(v),
        colorscale=[[0, "#9aa0a6"], [1, "#9aa0a6"]],
        opacity=0.12,
        showscale=False,
        hoverinfo="skip",
        name="template sphere",
    )
)

# One city is called out on its own so it can be found instantly among several thousand markers.
HIGHLIGHT_CITY = "Lisbon"
highlight = plot_df[plot_df["city"] == HIGHLIGHT_CITY]

for continent, group in plot_df[plot_df["city"] != HIGHLIGHT_CITY].groupby("continent", sort=False):
    for split, marker in (("train", "circle"), ("held-out", "diamond-open")):
        subset = group[group["split"] == split]
        if subset.empty:
            continue
        fig.add_trace(
            go.Scatter3d(
                x=subset["x"],
                y=subset["y_"],
                z=subset["z"],
                mode="markers",
                name=f"{continent} ({split})",
                legendgroup=continent,
                marker=dict(size=2, symbol=marker, color=CONTINENT_COLORS[continent], line=dict(width=0)),
                customdata=np.stack(
                    [subset["city"], subset["country"], subset["lat"], subset["lon"], subset["split"]], axis=1
                ),
                hovertemplate=(
                    "<b>%{customdata[0]}</b>, %{customdata[1]}<br>"
                    + continent
                    + " · %{customdata[4]}<br>"
                    "lat %{customdata[2]:.2f}, lon %{customdata[3]:.2f}<extra></extra>"
                ),
            )
        )

# Scatter3d offers no star glyph, so the call-out is a heavy "x" in ink rather than a series colour.
fig.add_trace(
    go.Scatter3d(
        x=highlight["x"],
        y=highlight["y_"],
        z=highlight["z"],
        mode="markers+text",
        name=HIGHLIGHT_CITY,
        marker=dict(size=5, symbol="x", color="#111827", line=dict(width=1, color="#ffffff")),
        text=highlight["city"],
        textposition="top center",
        textfont=dict(size=13, color="#111827"),
        customdata=np.stack(
            [highlight["city"], highlight["country"], highlight["lat"], highlight["lon"], highlight["split"]], axis=1
        ),
        hovertemplate=(
            "<b>%{customdata[0]}</b>, %{customdata[1]}<br>"
            "%{customdata[4]}<br>"
            "lat %{customdata[2]:.2f}, lon %{customdata[3]:.2f}<extra></extra>"
        ),
    )
)

fig.update_layout(
    title="Spherical SMDS embedding of city activations (aligned to the Earth frame)",
    template="plotly_white",
    height=820,
    margin=dict(l=0, r=0, t=60, b=0),
    # Small legend floated inside the plotting area, so the globe keeps the full width of the figure
    # instead of being squeezed into whatever a legend sitting beside it leaves over.
    legend=dict(
        itemsizing="constant",
        title=dict(text="Continent · split", font=dict(size=10)),
        font=dict(size=9),
        itemwidth=30,
        xanchor="right",
        x=1,
        yanchor="top",
        y=1,
        bgcolor="rgba(255,255,255,0.6)",
    ),
    scene=dict(
        xaxis_title="toward 0°N 0°E",
        yaxis_title="toward 0°N 90°E",
        zaxis_title="toward the North Pole",
        aspectmode="data",
        domain=dict(x=[0, 1], y=[0, 1]),
        camera=dict(eye=dict(x=1.05, y=1.05, z=0.85)),  # default eye ~1.25 each; closer = bigger globe
    ),
)
fig.show()
# Fixed categorical order -- hues are assigned per continent and never cycled.
CONTINENT_COLORS = {
    "Europe": "#2a78d6",
    "Africa": "#eb6834",
    "Asia": "#1baf7a",
    "North America": "#eda100",
    "South America": "#e87ba4",
    "Oceania": "#4a3aa7",
}

def fit_sphere(points: NDArray[np.float64]) -> tuple[NDArray[np.float64], float]:
    """Least-squares sphere through the points, as (centre, radius)."""
    design = np.c_[2 * points, np.ones(len(points))]
    solution, *_ = np.linalg.lstsq(design, (points**2).sum(axis=1), rcond=None)
    centre = solution[:3]
    radius = float(np.sqrt(solution[3] + centre @ centre))
    return centre, radius


def globe_coords(coordinates: NDArray[np.float64]) -> NDArray[np.float64]:
    """Map (lat, lon) in degrees onto the unit sphere, z pointing at the North Pole."""
    lat, lon = np.radians(coordinates[:, 0]), np.radians(coordinates[:, 1])
    return np.stack([np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], axis=1)


spherical = models["spherical"]

# The template's own sphere. Classical MDS centres the embedding on the centroid of the sampled
# cities, which is not the centre of the globe, so the fitted centre is offset from the origin.
centre, radius = fit_sphere(spherical.Y_)
proj = spherical.transform(X_scaled) - centre

# An SMDS embedding is only defined up to an isometry: rotations and reflections leave every pairwise
# distance untouched, so the raw axes point in arbitrary directions and the globe comes out mirrored.
# Procrustes onto the true coordinates of the *training* cities puts the cloud in the Earth frame --
# z through the North Pole, x through (0 N, 0 E) -- which is what makes the picture readable.
rotation, _ = orthogonal_procrustes(proj[idx_train], globe_coords(y[idx_train]))
proj = proj @ rotation

print(f"template sphere: fitted centre {np.round(centre, 3)}, radius {radius:.3f}")
print(f"alignment to the Earth frame: {'reflection + rotation' if np.linalg.det(rotation) < 0 else 'rotation'}")

plot_df = df.assign(
    x=proj[:, 0],
    y_=proj[:, 1],
    z=proj[:, 2],
    split=np.where(np.isin(np.arange(len(df)), idx_train), "train", "held-out"),
)

u, v = np.mgrid[0 : 2 * np.pi : 60j, 0 : np.pi : 30j]
fig = go.Figure(
    go.Surface(
        x=radius * np.cos(u) * np.sin(v),
        y=radius * np.sin(u) * np.sin(v),
        z=radius * np.cos(v),
        colorscale=[[0, "#9aa0a6"], [1, "#9aa0a6"]],
        opacity=0.12,
        showscale=False,
        hoverinfo="skip",
        name="template sphere",
    )
)

# One city is called out on its own so it can be found instantly among several thousand markers.
HIGHLIGHT_CITY = "Lisbon"
highlight = plot_df[plot_df["city"] == HIGHLIGHT_CITY]

for continent, group in plot_df[plot_df["city"] != HIGHLIGHT_CITY].groupby("continent", sort=False):
    for split, marker in (("train", "circle"), ("held-out", "diamond-open")):
        subset = group[group["split"] == split]
        if subset.empty:
            continue
        fig.add_trace(
            go.Scatter3d(
                x=subset["x"],
                y=subset["y_"],
                z=subset["z"],
                mode="markers",
                name=f"{continent} ({split})",
                legendgroup=continent,
                marker=dict(size=2, symbol=marker, color=CONTINENT_COLORS[continent], line=dict(width=0)),
                customdata=np.stack(
                    [subset["city"], subset["country"], subset["lat"], subset["lon"], subset["split"]], axis=1
                ),
                hovertemplate=(
                    "%{customdata[0]}, %{customdata[1]}
"
                    + continent
                    + " · %{customdata[4]}
"
                    "lat %{customdata[2]:.2f}, lon %{customdata[3]:.2f}"
                ),
            )
        )

# Scatter3d offers no star glyph, so the call-out is a heavy "x" in ink rather than a series colour.
fig.add_trace(
    go.Scatter3d(
        x=highlight["x"],
        y=highlight["y_"],
        z=highlight["z"],
        mode="markers+text",
        name=HIGHLIGHT_CITY,
        marker=dict(size=5, symbol="x", color="#111827", line=dict(width=1, color="#ffffff")),
        text=highlight["city"],
        textposition="top center",
        textfont=dict(size=13, color="#111827"),
        customdata=np.stack(
            [highlight["city"], highlight["country"], highlight["lat"], highlight["lon"], highlight["split"]], axis=1
        ),
        hovertemplate=(
            "%{customdata[0]}, %{customdata[1]}
"
            "%{customdata[4]}
"
            "lat %{customdata[2]:.2f}, lon %{customdata[3]:.2f}"
        ),
    )
)

fig.update_layout(
    title="Spherical SMDS embedding of city activations (aligned to the Earth frame)",
    template="plotly_white",
    height=820,
    margin=dict(l=0, r=0, t=60, b=0),
    # Small legend floated inside the plotting area, so the globe keeps the full width of the figure
    # instead of being squeezed into whatever a legend sitting beside it leaves over.
    legend=dict(
        itemsizing="constant",
        title=dict(text="Continent · split", font=dict(size=10)),
        font=dict(size=9),
        itemwidth=30,
        xanchor="right",
        x=1,
        yanchor="top",
        y=1,
        bgcolor="rgba(255,255,255,0.6)",
    ),
    scene=dict(
        xaxis_title="toward 0°N 0°E",
        yaxis_title="toward 0°N 90°E",
        zaxis_title="toward the North Pole",
        aspectmode="data",
        domain=dict(x=[0, 1], y=[0, 1]),
        camera=dict(eye=dict(x=1.05, y=1.05, z=0.85)),  # default eye ~1.25 each; closer = bigger globe
    ),
)
fig.show()

template sphere: fitted centre [0.085 0.278 0.379], radius 1.000
alignment to the Earth frame: rotation

Stress comparison¶

The figure rests on one split; the numbers should not, so each geometry is refitted across N_FOLDS continent-stratified folds and scored on the fold it did not see.

SupervisedMDS.score returns $1 - \text{stress}$, so the table reports stress directly — lower is better — as both a metric (scale-normalised) and a rank-only (non-metric) measure.

Only the held-out columns carry evidence: with far more features than cities, the ridge fit drives train stress to ~0 for every geometry. The two curved templates finish close together and well ahead of the flat one, which is the takeaway — the model wraps longitude.

In [31]:

Copied!





def stress_scorer(metric: str):
    """Scorer reporting stress directly: SMDS returns 1 - stress, and the pipeline's scaler is
    applied by hand because `Pipeline.score` does not forward the `metric` argument."""

    def scorer(estimator, X_eval: NDArray[np.float32], y_eval: NDArray[np.float64]) -> float:
        return 1 - estimator[-1].score(estimator[:-1].transform(X_eval), y_eval, metric=metric)

    return scorer


cv = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
folds = list(cv.split(X, df["continent"]))  # materialised so every geometry sees the same folds

scoring = {
    "stress": stress_scorer("scale_normalized_stress"),
    "nm_stress": stress_scorer("non_metric_stress"),
}

rows = []
for name, model in models.items():
    # The scaler sits inside the pipeline, so each fold refits the centring on its own training
    # cities. N_JOBS folds run at once -- see the memory note where it is defined.
    cv_results = cross_validate(
        make_pipeline(StandardScaler(), model),
        X,
        y,
        cv=folds,
        scoring=scoring,
        return_train_score=True,
        n_jobs=N_JOBS,
        pre_dispatch="n_jobs",  # queue one fold per worker; the default 2x doubles peak memory
    )

    rows.append(
        {
            "geometry": name,
            "n_components": model.parametrization_fitted_.n_components,
            "train stress": cv_results["train_stress"].mean(),
            "held-out stress": cv_results["test_stress"].mean(),
            "held-out stress (sd)": cv_results["test_stress"].std(),
            "held-out non-metric stress": cv_results["test_nm_stress"].mean(),
        }
    )

stress_df = pd.DataFrame(rows).set_index("geometry")

best = stress_df["held-out stress"].idxmin()
runner_up = stress_df["held-out stress"].drop(best).idxmin()
gap = stress_df.loc[runner_up, "held-out stress"] - stress_df.loc[best, "held-out stress"]
print(f"Best geometry over {N_FOLDS} folds: {best}")
print(f"Ahead of {runner_up} by {gap:.3f}, against a fold-to-fold sd of {stress_df.loc[best, 'held-out stress (sd)']:.3f}\n")

stress_df.style.format({col: "{:.3f}" for col in stress_df.columns if col != "n_components"}).background_gradient(
    cmap="RdYlGn_r", subset=["held-out stress", "held-out non-metric stress"]
)
def stress_scorer(metric: str):
    """Scorer reporting stress directly: SMDS returns 1 - stress, and the pipeline's scaler is
    applied by hand because `Pipeline.score` does not forward the `metric` argument."""

    def scorer(estimator, X_eval: NDArray[np.float32], y_eval: NDArray[np.float64]) -> float:
        return 1 - estimator[-1].score(estimator[:-1].transform(X_eval), y_eval, metric=metric)

    return scorer


cv = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
folds = list(cv.split(X, df["continent"]))  # materialised so every geometry sees the same folds

scoring = {
    "stress": stress_scorer("scale_normalized_stress"),
    "nm_stress": stress_scorer("non_metric_stress"),
}

rows = []
for name, model in models.items():
    # The scaler sits inside the pipeline, so each fold refits the centring on its own training
    # cities. N_JOBS folds run at once -- see the memory note where it is defined.
    cv_results = cross_validate(
        make_pipeline(StandardScaler(), model),
        X,
        y,
        cv=folds,
        scoring=scoring,
        return_train_score=True,
        n_jobs=N_JOBS,
        pre_dispatch="n_jobs",  # queue one fold per worker; the default 2x doubles peak memory
    )

    rows.append(
        {
            "geometry": name,
            "n_components": model.parametrization_fitted_.n_components,
            "train stress": cv_results["train_stress"].mean(),
            "held-out stress": cv_results["test_stress"].mean(),
            "held-out stress (sd)": cv_results["test_stress"].std(),
            "held-out non-metric stress": cv_results["test_nm_stress"].mean(),
        }
    )

stress_df = pd.DataFrame(rows).set_index("geometry")

best = stress_df["held-out stress"].idxmin()
runner_up = stress_df["held-out stress"].drop(best).idxmin()
gap = stress_df.loc[runner_up, "held-out stress"] - stress_df.loc[best, "held-out stress"]
print(f"Best geometry over {N_FOLDS} folds: {best}")
print(f"Ahead of {runner_up} by {gap:.3f}, against a fold-to-fold sd of {stress_df.loc[best, 'held-out stress (sd)']:.3f}\n")

stress_df.style.format({col: "{:.3f}" for col in stress_df.columns if col != "n_components"}).background_gradient(
    cmap="RdYlGn_r", subset=["held-out stress", "held-out non-metric stress"]
)

Best geometry over 20 folds: cylindrical
Ahead of spherical by 0.004, against a fold-to-fold sd of 0.013

Out[31]:

	n_components	train stress	held-out stress	held-out stress (sd)	held-out non-metric stress
geometry
spherical	3	0.156	0.262	0.013	0.059
flat (linear)	3	0.192	0.316	0.021	0.089
cylindrical	3	0.153	0.258	0.013	0.056