欠損値ありデータへの協調フィルタリングには素のNMFではなく欠損値用NMFを使おう

import polars as pl


df = pl.read_csv("/content/ml-100k/u.data", sep="\t", has_header=False)
df.columns = ["user id", "item id", "rating", "timestamp"]

df_ranking = df.pivot(values="rating", index="user id", columns="item id").fill_null(0)
ranking_matrix = df_ranking[:, 1:].to_numpy()  # Rに対応

import numpy as np
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import mean_squared_error
from scipy.sparse import csr_matrix
from tqdm import tqdm

# ユーザの割合をなるべく均一に分けたいので、層化k分割交差検証を使う
user_ids = df_ranking["user id"].to_numpy()
row_idx, col_idx = ranking_matrix.nonzero()
user_labels = [user_ids[idx] for idx in row_idx]  # 基準となるユーザIDを抽出

n_iteration = 50  # NMFを実行することになる最大圧縮次元数-1
n_splits = 5  # 層化k分割交差検証での分割数
train_losses = np.zeros((n_iteration, n_splits))
val_losses = np.zeros((n_iteration, n_splits))
val_means = np.zeros((n_iteration, n_splits))


skf = StratifiedKFold(n_splits=n_splits, random_state=3, shuffle=True)
for i in tqdm(range(n_iteration)):
  train_loss = [None] * n_splits
  val_loss = [None] * n_splits
  val_mean = [None] * n_splits
  np.random.seed(3)
  for j, (train_index, test_index) in enumerate(skf.split(row_idx, user_labels)):
    train_ranking_matrix = np.zeros_like(ranking_matrix)  # 0埋めとWNMF用なので平均埋めの場合はこちらをコメント化する
    # train_ranking_matrix = np.full_like(ranking_matrix, ranking_matrix[row_idx[train_index], col_idx[train_index]].mean())  # 平均埋めの場合はこちらをコメントアウト
    train_ranking_matrix[row_idx[train_index], col_idx[train_index]] = ranking_matrix[row_idx[train_index], col_idx[train_index]]
    nmf = NMF(n_components=(i + 1))  # WNMFを使う場合はNMFをWNMFに置き換える
    P, Q = nmf(train_ranking_matrix)
    ranking_matrix_reconstruction = P @ Q
    train_pred = ranking_matrix_reconstruction[row_idx[train_index], col_idx[train_index]]
    val_pred = ranking_matrix_reconstruction[row_idx[test_index], col_idx[test_index]]
    train_true = ranking_matrix[row_idx[train_index], col_idx[train_index]]
    val_true = ranking_matrix[row_idx[test_index], col_idx[test_index]]
    train_loss[j] = mean_squared_error(train_true, train_pred, squared=False)
    val_loss[j] = mean_squared_error(val_true, val_pred, squared=False)
    val_mean[j] = val_pred.mean()
  train_losses[i] = np.asarray(train_loss)
  val_losses[i] = np.asarray(val_loss)
  val_means[i] = np.asarray(val_mean)

from typing import Tuple

import numpy as np


class NMF:
  def __init__(self, n_components: int = 100, maxiter: int = 100, **kwargs):
    self.n_components = n_components  # 圧縮次元数
    self.maxiter = maxiter  # 最大更新数

  def __call__(self, R: np.ndarray) -> Tuple[np.ndarray]:
    m, n = R.shape
    P = np.random.uniform(low=0, high=1, size=(m, self.n_components))
    Q = np.random.uniform(low=0, high=1, size=(self.n_components, n))
    for i in range(self.maxiter):
      P, Q = self.multiplicative_update_rule(R, P, Q)
    return P, Q

  def multiplicative_update_rule(self, R: np.ndarray, P: np.ndarray, Q: np.ndarray) -> Tuple[np.ndarray]:
    """ユークリッド距離によるNMFにおいて乗法更新則による更新後のPとQを求める
    """
    P_denominator = P @ Q @ Q.T  # Rに対するmaskは変わらないので除いてある
    P_denominator[P_denominator == 0] = 1.0E-7  # 0除算対策
    P = P * (R @ Q.T) / P_denominator
    Q_denominator = P.T @ P @ Q
    Q_denominator[Q_denominator == 0] = 1.0E-7
    Q = Q * (P.T @  R) / Q_denominator
    return P, Q


class WNMF(NMF):
  def __init__(self, n_components: int = 100, maxiter: int = 100, **kwargs):
    super().__init__(n_components=n_components, maxiter=maxiter, **kwargs)

  def multiplicative_update_rule(self, R: np.ndarray, P: np.ndarray, Q: np.ndarray) -> Tuple[np.ndarray]:
    """ユークリッド距離によるWNMFにおいて乗法更新則による更新後のPとQを求める

    Rは欠損値を既に0で埋めていることを想定している
    """
    P_denominator = P @ Q @ Q.T
    P_denominator[P_denominator == 0] = 1.0E-7  # 0除算対策
    P = P * (R @ Q.T) / P_denominator
    Q_denominator = P.T @ P @ Q
    Q_denominator[Q_denominator == 0] = 1.0E-7
    Q = Q * (P.T @  R) / Q_denominator
    return P, Q