Factor analysis of mixed data

from __future__ import annotations

import prince

dataset = prince.datasets.load_beers().head(1000)
dataset.head()

famd = prince.FAMD(
    n_components=2,
    n_iter=3,
    rescale_with_mean=True,
    rescale_with_std=False,
    copy=True,
    check_input=True,
    random_state=42,
    engine="sklearn",
    handle_unknown="error",  # same parameter as sklearn.preprocessing.OneHotEncoder
)
famd = famd.fit(dataset)

famd.eigenvalues_summary

famd.row_coordinates(dataset).head()

famd.column_coordinates_

famd.plot(dataset, x_component=0, y_component=1)

(famd.row_contributions_.sort_values(0, ascending=False).head(5).style.format("{:.3%}"))

famd.column_contributions_.style.format("{:.0%}")

import pandas as pd

wiki = pd.DataFrame(
    {
        "k1": [2.0, 5.0, 3.0, 4.0, 1.0, 6.0],
        "k2": [4.5, 4.5, 1.0, 1.0, 1.0, 1.0],
        "k3": [4.0, 4.0, 2.0, 2.0, 1.0, 2.0],
        "q1": ["A", "C", "B", "B", "A", "C"],
        "q2": ["B", "B", "B", "B", "A", "A"],
        "q3": ["C", "C", "B", "B", "A", "A"],
    },
    index=[f"i{i}" for i in range(1, 7)],
)
wiki

wiki_famd = prince.FAMD(n_components=5, engine="scipy")
wiki_famd = wiki_famd.fit(wiki)
wiki_famd.eigenvalues_summary

import altair as alt

rc = wiki_famd.row_coordinates(wiki)
rc.columns = [f"F{c + 1}" for c in rc.columns]
rc["label"] = rc.index

alt.Chart(rc.reset_index(drop=True)).mark_text(fontSize=14).encode(
    x=alt.X("F1:Q", title="F1"),
    y=alt.Y("F2:Q", title="F2"),
    text="label:N",
).properties(title="Figure 1 — Individuals", width=400, height=300)

num_cols = ["k1", "k2", "k3"]
cat_cols = ["q1", "q2", "q3"]

# FactoMineR's var$coord: r² for numerical, η² for categorical (one row per variable).
cc = wiki_famd.variable_coordinates_.iloc[:, :2].copy()
cc.columns = ["F1", "F2"]
cc["variable"] = cc.index
cc["type"] = ["quanti"] * len(num_cols) + ["quali"] * len(cat_cols)

points = (
    alt.Chart(cc.reset_index(drop=True))
    .mark_text(fontSize=14)
    .encode(
        x=alt.X("F1:Q", title="F1 (r² / η²)", scale=alt.Scale(domain=[0, 1.05])),
        y=alt.Y("F2:Q", title="F2 (r² / η²)", scale=alt.Scale(domain=[0, 1.05])),
        text="variable:N",
        color=alt.Color("type:N", legend=alt.Legend(title="Variable type")),
    )
)

points.properties(title="Figure 2 — Relationship square", width=400, height=400)

import numpy as np

num_cols = ["k1", "k2", "k3"]
cat_cols = ["q1", "q2", "q3"]

corr_circle = wiki_famd.column_correlations.loc[num_cols].iloc[:, :2].copy()
corr_circle.columns = ["F1", "F2"]
corr_circle["variable"] = corr_circle.index

# Arrows from origin to each variable
arrows = pd.DataFrame(
    [
        {"x": 0, "y": 0, "x2": row["F1"], "y2": row["F2"], "variable": row["variable"]}
        for _, row in corr_circle.iterrows()
    ]
)

circle_theta = np.linspace(0, 2 * np.pi, 100)
unit_circle = pd.DataFrame(
    {"x": np.cos(circle_theta), "y": np.sin(circle_theta), "order": range(100)}
)

c_circle = (
    alt.Chart(unit_circle)
    .mark_line(color="gray", strokeDash=[4, 4])
    .encode(
        x=alt.X("x:Q", title="F1", scale=alt.Scale(domain=[-1.2, 1.2])),
        y=alt.Y("y:Q", title="F2", scale=alt.Scale(domain=[-1.2, 1.2])),
        order="order:O",
    )
)
c_arrows = (
    alt.Chart(arrows)
    .mark_rule(strokeWidth=2)
    .encode(
        x="x:Q",
        y="y:Q",
        x2="x2:Q",
        y2="y2:Q",
        color="variable:N",
    )
)
c_labels = (
    alt.Chart(corr_circle.reset_index(drop=True))
    .mark_text(fontSize=14, dx=15)
    .encode(x="F1:Q", y="F2:Q", text="variable:N")
)

(c_circle + c_arrows + c_labels).properties(
    title="Figure 3 — Correlation circle", width=400, height=400
)

# Category coordinates are the mean of row coordinates for each category
cat_coords = []
for q in cat_cols:
    for cat in sorted(wiki[q].unique()):
        mean_rc = rc.loc[wiki[q] == cat, ["F1", "F2"]].mean()
        cat_coords.append({"F1": mean_rc["F1"], "F2": mean_rc["F2"], "label": f"{q}={cat}"})

cat_df = pd.DataFrame(cat_coords)

alt.Chart(cat_df).mark_text(fontSize=14).encode(
    x=alt.X("F1:Q", title="F1"),
    y=alt.Y("F2:Q", title="F2"),
    text="label:N",
).properties(title="Figure 4 — Categories", width=400, height=300)