test_interactions.py

#  Copyright 2022-2024 MTS (Mobile Telesystems)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

# pylint: disable=attribute-defined-outside-init

import typing as tp
from datetime import datetime

import numpy as np
import pandas as pd
import pytest
from pytest_subtests import SubTests
from scipy import sparse

from rectools import Columns
from rectools.dataset import IdMap, Interactions
from tests.testing_utils import assert_sparse_matrix_equal


class TestInteractions:
    def setup_method(self) -> None:
        self.df = pd.DataFrame(
            {
                Columns.User: [1, 2, 1, 1],
                Columns.Item: [0, 1, 0, 1],
                Columns.Weight: [5, 7.0, 4, 1],
                Columns.Datetime: [datetime(2021, 9, 8)] * 4,
                "extra_col": [1, 2, 3, 4],
            }
        )
        self.raw_df = pd.DataFrame(
            {
                Columns.User: ["u1", "u2", "u1", "u1"],
                Columns.Item: ["i1", "i2", "i1", "i2"],
                Columns.Weight: [5, 7, 4, 1],
                Columns.Datetime: ["2021-09-08"] * 4,
                "extra_col": [1, 2, 3, 4],
            }
        )

    def test_creation(self) -> None:
        interactions = Interactions(self.df)
        pd.testing.assert_frame_equal(interactions.df, self.df)

    def test_missing_columns_validation(self, subtests: SubTests) -> None:
        for col in self.df.columns:
            with subtests.test(f"drop {col} column"):
                if col != "extra_col":
                    with pytest.raises(KeyError):
                        Interactions(self.df.drop(columns=col))

    @pytest.mark.parametrize("column", (Columns.User, Columns.Item))
    def test_types_validation(self, column: str) -> None:
        with pytest.raises(TypeError):
            Interactions(self.df.astype({column: float}))

    @pytest.mark.parametrize("column", (Columns.User, Columns.Item))
    def test_positivity_validation(self, column: str) -> None:
        with pytest.raises(ValueError):
            self.df.at[0, column] = -1
            Interactions(self.df)

    @pytest.mark.parametrize("keep_extra_cols", (True, False))
    def test_from_raw_creation(self, keep_extra_cols: bool) -> None:
        raw_df = self.raw_df
        user_id_map = IdMap(np.array(["u0", "u1", "u2"]))
        item_id_map = IdMap.from_values(["i1", "i2"])
        interactions = Interactions.from_raw(raw_df, user_id_map, item_id_map, keep_extra_cols=keep_extra_cols)
        excepted = self.df
        if not keep_extra_cols:
            excepted.drop(columns="extra_col", inplace=True)
        pd.testing.assert_frame_equal(interactions.df, excepted)

    @pytest.mark.parametrize(
        "with_weights,expected_data",
        (
            (False, [1, 1, 1, 1]),
            (True, [5, 7, 4, 1]),
        ),
    )
    def test_getting_user_item_matrix(self, with_weights: bool, expected_data: tp.List[float]) -> None:
        interactions = Interactions(self.df)
        matrix = interactions.get_user_item_matrix(with_weights)
        expected = sparse.csr_matrix((expected_data, (self.df[Columns.User].values, self.df[Columns.Item].values)))
        assert_sparse_matrix_equal(matrix, expected)

    def test_raises_when_weight_not_numeric(self) -> None:
        df = self.df
        df.loc[1, Columns.Weight] = "w"
        with pytest.raises(TypeError) as e:
            Interactions.from_raw(df, IdMap.from_values(df[Columns.User]), IdMap.from_values(df[Columns.Item]))
        err_text = e.value.args[0]
        assert Columns.Weight in err_text.lower()

    def test_raises_when_datetime_type_incorrect(self) -> None:
        df = self.df
        df.loc[1, Columns.Datetime] = "dt"
        with pytest.raises(TypeError) as e:
            Interactions.from_raw(df, IdMap.from_values(df[Columns.User]), IdMap.from_values(df[Columns.Item]))
        err_text = e.value.args[0]
        assert Columns.Datetime in err_text.lower()

    @pytest.mark.parametrize("include_weight", (True, False))
    @pytest.mark.parametrize("include_datetime", (True, False))
    @pytest.mark.parametrize("include_extra_cols", (True, False))
    def test_to_external(self, include_weight: bool, include_datetime: bool, include_extra_cols: bool) -> None:
        user_id_map = IdMap(np.array([10, 20, 30]))
        item_id_map = IdMap(np.array(["i1", "i2"]))
        interactions = Interactions(self.df)

        actual = interactions.to_external(
            user_id_map, item_id_map, include_weight, include_datetime, include_extra_cols
        )
        expected = pd.DataFrame(
            [
                [20, "i1"],
                [30, "i2"],
                [20, "i1"],
                [20, "i2"],
            ],
            columns=Columns.UserItem,
        )
        if include_weight:
            expected[Columns.Weight] = self.df[Columns.Weight]
        if include_datetime:
            expected[Columns.Datetime] = self.df[Columns.Datetime]
        if include_extra_cols:
            expected["extra_col"] = self.df["extra_col"]

        pd.testing.assert_frame_equal(actual, expected)

    def test_to_external_empty(self) -> None:
        user_id_map = IdMap(np.array([10, 20, 30]))
        item_id_map = IdMap(np.array(["i1", "i2"]))
        interactions = Interactions(self.df.iloc[:0])

        actual = interactions.to_external(user_id_map, item_id_map, include_extra_cols=False)
        expected = pd.DataFrame(
            [],
            columns=Columns.Interactions,
        )
        expected = expected.astype(
            {
                Columns.User: np.int64,
                Columns.Item: "object",
                Columns.Weight: np.float64,
                Columns.Datetime: "datetime64[ns]",
            }
        )
        pd.testing.assert_frame_equal(actual, expected, check_index_type=False)

    def test_to_external_with_missing_ids(self) -> None:
        user_id_map = IdMap(np.array([10, 20, 30]))
        item_id_map = IdMap(np.array(["i1"]))
        interactions = Interactions(self.df)

        with pytest.raises(KeyError):
            interactions.to_external(user_id_map, item_id_map)