"""Test cases for PyKEEN."""
import inspect
import itertools
import logging
import os
import pathlib
import tempfile
import timeit
import traceback
import unittest
from abc import ABC, abstractmethod
from collections import ChainMap, Counter
from collections.abc import Collection, Iterable, Mapping, MutableMapping, Sequence
from typing import (
Any,
Callable,
ClassVar,
Optional,
TypeVar,
Union,
)
from unittest.case import SkipTest
from unittest.mock import Mock, patch
import numpy
import numpy.random
import pandas
import pytest
import torch
import torch.utils.data
import unittest_templates
from class_resolver import HintOrType
from click.testing import CliRunner, Result
from docdata import get_docdata
from torch import optim
from torch.nn import functional
from torch.optim import SGD, Adagrad
import pykeen.evaluation.evaluation_loop
import pykeen.models
import pykeen.nn.combination
import pykeen.nn.message_passing
import pykeen.nn.node_piece
import pykeen.nn.representation
import pykeen.nn.text
import pykeen.nn.weighting
import pykeen.predict
from pykeen.datasets import Nations
from pykeen.datasets.base import LazyDataset
from pykeen.datasets.ea.combination import GraphPairCombinator
from pykeen.datasets.kinships import KINSHIPS_TRAIN_PATH
from pykeen.datasets.mocks import create_inductive_dataset
from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH
from pykeen.evaluation import Evaluator, MetricResults, evaluator_resolver
from pykeen.losses import Loss, PairwiseLoss, PointwiseLoss, SetwiseLoss, UnsupportedLabelSmoothingError, loss_resolver
from pykeen.metrics import rank_based_metric_resolver
from pykeen.metrics.ranking import (
DerivedRankBasedMetric,
NoClosedFormError,
RankBasedMetric,
generate_num_candidates_and_ranks,
)
from pykeen.models import RESCAL, ERModel, Model, TransE
from pykeen.models.cli import build_cli_from_cls
from pykeen.models.meta.filtered import CooccurrenceFilteredModel
from pykeen.models.mocks import FixedModel
from pykeen.nn.modules import DistMultInteraction, FunctionalInteraction, Interaction
from pykeen.nn.representation import Representation
from pykeen.nn.utils import adjacency_tensor_to_stacked_matrix
from pykeen.optimizers import optimizer_resolver
from pykeen.pipeline import pipeline
from pykeen.regularizers import LpRegularizer, Regularizer
from pykeen.stoppers.early_stopping import EarlyStopper
from pykeen.trackers import ResultTracker
from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingCallback, TrainingLoop
from pykeen.triples import Instances, TriplesFactory, generation
from pykeen.triples.instances import BaseBatchedSLCWAInstances, SLCWABatch
from pykeen.triples.splitting import Cleaner, Splitter
from pykeen.triples.triples_factory import CoreTriplesFactory
from pykeen.triples.utils import get_entities
from pykeen.typing import (
EA_SIDE_LEFT,
EA_SIDE_RIGHT,
LABEL_HEAD,
LABEL_TAIL,
RANK_REALISTIC,
SIDE_BOTH,
TRAINING,
HeadRepresentation,
InductiveMode,
Initializer,
MappedTriples,
RelationRepresentation,
TailRepresentation,
Target,
)
from pykeen.utils import (
all_in_bounds,
get_batchnorm_modules,
getattr_or_docdata,
is_triple_tensor_subset,
resolve_device,
set_random_seed,
triple_tensor_to_set,
unpack_singletons,
)
from tests.constants import EPSILON
from tests.mocks import MockEvaluator
from tests.utils import needs_packages, rand
T = TypeVar("T")
logger = logging.getLogger(__name__)
class GenericTestCase(unittest_templates.GenericTestCase[T]):
"""Generic tests."""
generator: torch.Generator
def pre_setup_hook(self) -> None:
"""Instantiate a generator for usage in the test case."""
self.generator = set_random_seed(seed=42)[1]
class DatasetTestCase(unittest.TestCase):
"""A test case for quickly defining common tests for datasets."""
#: The expected number of entities
exp_num_entities: ClassVar[int]
#: The expected number of relations
exp_num_relations: ClassVar[int]
#: The expected number of triples
exp_num_triples: ClassVar[int]
#: The tolerance on expected number of triples, for randomized situations
exp_num_triples_tolerance: ClassVar[Optional[int]] = None
#: The dataset to test
dataset_cls: ClassVar[type[LazyDataset]]
#: The instantiated dataset
dataset: LazyDataset
#: Should the validation be assumed to have been loaded with train/test?
autoloaded_validation: ClassVar[bool] = False
def test_dataset(self):
"""Generic test for datasets."""
self.assertIsInstance(self.dataset, LazyDataset)
# Not loaded
self.assertIsNone(self.dataset._training)
self.assertIsNone(self.dataset._testing)
self.assertIsNone(self.dataset._validation)
self.assertFalse(self.dataset._loaded)
self.assertFalse(self.dataset._loaded_validation)
# Load
self.dataset._load()
self.assertIsInstance(self.dataset.training, TriplesFactory)
self.assertIsInstance(self.dataset.testing, TriplesFactory)
self.assertTrue(self.dataset._loaded)
if self.autoloaded_validation:
self.assertTrue(self.dataset._loaded_validation)
else:
self.assertFalse(self.dataset._loaded_validation)
self.dataset._load_validation()
self.assertIsInstance(self.dataset.validation, TriplesFactory)
self.assertIsNotNone(self.dataset._training)
self.assertIsNotNone(self.dataset._testing)
self.assertIsNotNone(self.dataset._validation)
self.assertTrue(self.dataset._loaded)
self.assertTrue(self.dataset._loaded_validation)
self.assertEqual(self.dataset.num_entities, self.exp_num_entities)
self.assertEqual(self.dataset.num_relations, self.exp_num_relations)
num_triples = sum(
triples_factory.num_triples
for triples_factory in (self.dataset._training, self.dataset._testing, self.dataset._validation)
)
if self.exp_num_triples_tolerance is None:
self.assertEqual(self.exp_num_triples, num_triples)
else:
self.assertAlmostEqual(self.exp_num_triples, num_triples, delta=self.exp_num_triples_tolerance)
# Test caching
start = timeit.default_timer()
_ = self.dataset.training
end = timeit.default_timer()
# assert (end - start) < 1.0e-02
self.assertAlmostEqual(start, end, delta=1.0e-02, msg="Caching should have made this operation fast")
# Test consistency of training / validation / testing mapping
training = self.dataset.training
for part, factory in self.dataset.factory_dict.items():
if not isinstance(factory, TriplesFactory):
logger.warning("Skipping mapping consistency checks since triples factory does not provide mappings.")
continue
if part == "training":
continue
assert training.entity_to_id == factory.entity_to_id
assert training.num_entities == factory.num_entities
assert training.relation_to_id == factory.relation_to_id
assert training.num_relations == factory.num_relations
class LocalDatasetTestCase(DatasetTestCase):
"""A test case for datasets that don't need a cache directory."""
def setUp(self):
"""Set up the test case."""
self.dataset = self.dataset_cls()
class CachedDatasetCase(DatasetTestCase):
"""A test case for datasets that need a cache directory."""
#: The directory, if there is caching
directory: Optional[tempfile.TemporaryDirectory]
def setUp(self):
"""Set up the test with a temporary cache directory."""
self.directory = tempfile.TemporaryDirectory()
self.dataset = self.dataset_cls(cache_root=self.directory.name)
def tearDown(self) -> None:
"""Tear down the test case by cleaning up the temporary cache directory."""
self.directory.cleanup()
def iter_from_space(key: str, space: Mapping[str, Any]) -> Iterable[MutableMapping[str, Any]]:
"""Iterate over some configurations for a single hyperparameter."""
typ = space["type"]
if typ == "categorical":
for choice in space["choices"]:
yield {key: choice}
elif typ in (float, int):
yield {key: space["low"]}
yield {key: space["high"]}
else:
raise NotImplementedError(typ)
def iter_hpo_configs(hpo_default: Mapping[str, Mapping[str, Any]]) -> Iterable[Mapping[str, Any]]:
"""Iterate over some representative configurations from the HPO default."""
for combination in itertools.product(
*(iter_from_space(key=key, space=space) for key, space in hpo_default.items())
):
yield ChainMap(*combination)
class LossTestCase(GenericTestCase[Loss]):
"""Base unittest for loss functions."""
#: The batch size
batch_size: ClassVar[int] = 3
#: The number of negatives per positive for sLCWA training loop.
num_neg_per_pos: ClassVar[int] = 7
#: The number of entities LCWA training loop / label smoothing.
num_entities: ClassVar[int] = 7
def _check_loss_value(self, loss_value: torch.FloatTensor) -> None:
"""Check loss value dimensionality, and ability for backward."""
# test reduction
self.assertEqual(0, loss_value.ndim)
# test finite loss value
self.assertTrue(torch.isfinite(loss_value))
# Test backward
loss_value.backward()
def help_test_process_slcwa_scores(
self,
positive_scores: torch.FloatTensor,
negative_scores: torch.FloatTensor,
batch_filter: Optional[torch.BoolTensor] = None,
):
"""Help test processing scores from SLCWA training loop."""
loss_value = self.instance.process_slcwa_scores(
positive_scores=positive_scores,
negative_scores=negative_scores,
label_smoothing=None,
batch_filter=batch_filter,
num_entities=self.num_entities,
)
self._check_loss_value(loss_value=loss_value)
def test_process_slcwa_scores(self):
"""Test processing scores from SLCWA training loop."""
positive_scores = torch.rand(self.batch_size, 1, requires_grad=True)
negative_scores = torch.rand(self.batch_size, self.num_neg_per_pos, requires_grad=True)
self.help_test_process_slcwa_scores(positive_scores=positive_scores, negative_scores=negative_scores)
def test_process_slcwa_scores_filtered(self):
"""Test processing scores from SLCWA training loop with filtering."""
positive_scores = torch.rand(self.batch_size, 1, requires_grad=True)
negative_scores = torch.rand(self.batch_size, self.num_neg_per_pos, requires_grad=True)
batch_filter = torch.rand(self.batch_size, self.num_neg_per_pos) < 0.5
self.help_test_process_slcwa_scores(
positive_scores=positive_scores,
negative_scores=negative_scores[batch_filter],
batch_filter=batch_filter,
)
def test_process_lcwa_scores(self):
"""Test processing scores from LCWA training loop without smoothing."""
self.help_test_process_lcwa_scores(label_smoothing=None)
def test_process_lcwa_scores_smooth(self):
"""Test processing scores from LCWA training loop with smoothing."""
try:
self.help_test_process_lcwa_scores(label_smoothing=0.01)
except UnsupportedLabelSmoothingError as error:
raise SkipTest from error
def help_test_process_lcwa_scores(self, label_smoothing):
"""Help test processing scores from LCWA training loop."""
predictions = torch.rand(self.batch_size, self.num_entities, requires_grad=True)
labels = (torch.rand(self.batch_size, self.num_entities, requires_grad=True) > 0.8).float()
loss_value = self.instance.process_lcwa_scores(
predictions=predictions,
labels=labels,
label_smoothing=label_smoothing,
num_entities=self.num_entities,
)
self._check_loss_value(loss_value=loss_value)
def test_optimization_direction_lcwa(self):
"""Test whether the loss leads to increasing positive scores, and decreasing negative scores."""
labels = torch.as_tensor(data=[0, 1], dtype=torch.get_default_dtype()).view(1, -1)
predictions = torch.zeros(1, 2, requires_grad=True)
optimizer = optimizer_resolver.make(query=None, params=[predictions])
for _ in range(10):
optimizer.zero_grad()
loss = self.instance.process_lcwa_scores(predictions=predictions, labels=labels)
loss.backward()
optimizer.step()
# negative scores decreased compared to positive ones
assert predictions[0, 0] < predictions[0, 1] - 1.0e-06
def test_optimization_direction_slcwa(self):
"""Test whether the loss leads to increasing positive scores, and decreasing negative scores."""
positive_scores = torch.zeros(self.batch_size, 1, requires_grad=True)
negative_scores = torch.zeros(self.batch_size, self.num_neg_per_pos, requires_grad=True)
optimizer = optimizer_resolver.make(query=None, params=[positive_scores, negative_scores])
for _ in range(10):
optimizer.zero_grad()
loss = self.instance.process_slcwa_scores(
positive_scores=positive_scores,
negative_scores=negative_scores,
)
loss.backward()
optimizer.step()
# negative scores decreased compared to positive ones
assert (negative_scores < positive_scores.unsqueeze(dim=1) - 1.0e-06).all()
def test_hpo_defaults(self):
"""Test hpo defaults."""
signature = inspect.signature(self.cls.__init__)
# check for invalid keys
invalid_keys = set(self.cls.hpo_default.keys()).difference(signature.parameters)
assert not invalid_keys
# check that each parameter without a default occurs
required_parameters = {
parameter_name
for parameter_name, parameter in signature.parameters.items()
if parameter.default == inspect.Parameter.empty
and parameter.kind not in {inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL}
}.difference({"self"})
missing_required = required_parameters.difference(self.cls.hpo_default.keys())
assert not missing_required
# try to instantiate loss for some configurations in the HPO search space
for config in iter_hpo_configs(self.cls.hpo_default):
assert loss_resolver.make(query=self.cls, pos_kwargs=config)
class PointwiseLossTestCase(LossTestCase):
"""Base unit test for label-based losses."""
#: The number of entities.
num_entities: int = 17
def test_type(self):
"""Test the loss is the right type."""
self.assertIsInstance(self.instance, PointwiseLoss)
def test_label_loss(self):
"""Test ``forward(logits, labels)``."""
logits = torch.rand(self.batch_size, self.num_entities, requires_grad=True)
labels = functional.normalize(torch.rand(self.batch_size, self.num_entities, requires_grad=False), p=1, dim=-1)
loss_value = self.instance(
logits,
labels,
)
self._check_loss_value(loss_value)
class PairwiseLossTestCase(LossTestCase):
"""Base unit test for pair-wise losses."""
#: The number of negative samples
num_negatives: int = 5
def test_type(self):
"""Test the loss is the right type."""
self.assertIsInstance(self.instance, PairwiseLoss)
def test_pair_loss(self):
"""Test ``forward(pos_scores, neg_scores)``."""
pos_scores = torch.rand(self.batch_size, 1, requires_grad=True)
neg_scores = torch.rand(self.batch_size, self.num_negatives, requires_grad=True)
loss_value = self.instance(
pos_scores,
neg_scores,
)
self._check_loss_value(loss_value)
class GMRLTestCase(PairwiseLossTestCase):
"""Tests for generalized margin ranking loss."""
def test_label_smoothing_raise(self):
"""Test errors are raised if label smoothing is given."""
with self.assertRaises(UnsupportedLabelSmoothingError):
self.instance.process_lcwa_scores(..., ..., label_smoothing=5)
with self.assertRaises(UnsupportedLabelSmoothingError):
self.instance.process_lcwa_scores(..., ..., label_smoothing=5)
class SetwiseLossTestCase(LossTestCase):
"""Unit tests for setwise losses."""
#: The number of entities.
num_entities: int = 13
def test_type(self):
"""Test the loss is the right type."""
self.assertIsInstance(self.instance, SetwiseLoss)
class InteractionTestCase(
GenericTestCase[Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation]],
ABC,
):
"""Generic test for interaction functions."""
dim: int = 2
batch_size: int = 3
num_relations: int = 5
num_entities: int = 7
dtype: torch.dtype = torch.get_default_dtype()
# the relative tolerance for checking close results, cf. torch.allclose
rtol: float = 1.0e-5
# the absolute tolerance for checking close results, cf. torch.allclose
atol: float = 1.0e-8
shape_kwargs = dict()
def post_instantiation_hook(self) -> None:
"""Initialize parameters."""
self.instance.reset_parameters()
def _get_hrt(
self,
*shapes: tuple[int, ...],
):
shape_kwargs = dict(self.shape_kwargs)
shape_kwargs.setdefault("d", self.dim)
result = tuple(
tuple(
torch.rand(
size=tuple(prefix_shape) + tuple(shape_kwargs[dim] for dim in weight_shape),
requires_grad=True,
dtype=self.dtype,
)
for weight_shape in weight_shapes
)
for prefix_shape, weight_shapes in zip(
shapes,
[self.instance.entity_shape, self.instance.relation_shape, self.instance.tail_entity_shape],
)
)
return unpack_singletons(*result)
def _check_scores(self, scores: torch.FloatTensor, exp_shape: tuple[int, ...]):
"""Check shape, dtype and gradients of scores."""
assert torch.is_tensor(scores)
assert scores.dtype == torch.float32
assert scores.ndimension() == len(exp_shape)
assert scores.shape == exp_shape
assert scores.requires_grad
self._additional_score_checks(scores)
def _additional_score_checks(self, scores):
"""Additional checks for scores."""
@property
def _score_batch_sizes(self) -> Iterable[int]:
"""Return the list of batch sizes to test."""
if get_batchnorm_modules(self.instance):
return [self.batch_size]
return [1, self.batch_size]
def test_score_hrt(self):
"""Test score_hrt."""
for batch_size in self._score_batch_sizes:
h, r, t = self._get_hrt(
(batch_size,),
(batch_size,),
(batch_size,),
)
scores = self.instance.score_hrt(h=h, r=r, t=t)
self._check_scores(scores=scores, exp_shape=(batch_size, 1))
def test_score_h(self):
"""Test score_h."""
for batch_size in self._score_batch_sizes:
h, r, t = self._get_hrt(
(self.num_entities,),
(batch_size,),
(batch_size,),
)
scores = self.instance.score_h(all_entities=h, r=r, t=t)
self._check_scores(scores=scores, exp_shape=(batch_size, self.num_entities))
def test_score_h_slicing(self):
"""Test score_h with slicing."""
#: The equivalence for models with batch norm only holds in evaluation mode
self.instance.eval()
h, r, t = self._get_hrt(
(self.num_entities,),
(self.batch_size,),
(self.batch_size,),
)
scores = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=self.num_entities // 2 + 1)
scores_no_slice = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=None)
self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice)
def test_score_r(self):
"""Test score_r."""
for batch_size in self._score_batch_sizes:
h, r, t = self._get_hrt(
(batch_size,),
(self.num_relations,),
(batch_size,),
)
scores = self.instance.score_r(h=h, all_relations=r, t=t)
if len(self.cls.relation_shape) == 0:
exp_shape = (batch_size, 1)
else:
exp_shape = (batch_size, self.num_relations)
self._check_scores(scores=scores, exp_shape=exp_shape)
def test_score_r_slicing(self):
"""Test score_r with slicing."""
if len(self.cls.relation_shape) == 0:
raise unittest.SkipTest("No use in slicing relations for models without relation information.")
#: The equivalence for models with batch norm only holds in evaluation mode
self.instance.eval()
h, r, t = self._get_hrt(
(self.batch_size,),
(self.num_relations,),
(self.batch_size,),
)
scores = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=self.num_relations // 2 + 1)
scores_no_slice = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=None)
self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice)
def test_score_t(self):
"""Test score_t."""
for batch_size in self._score_batch_sizes:
h, r, t = self._get_hrt(
(batch_size,),
(batch_size,),
(self.num_entities,),
)
scores = self.instance.score_t(h=h, r=r, all_entities=t)
self._check_scores(scores=scores, exp_shape=(batch_size, self.num_entities))
def test_score_t_slicing(self):
"""Test score_t with slicing."""
#: The equivalence for models with batch norm only holds in evaluation mode
self.instance.eval()
h, r, t = self._get_hrt(
(self.batch_size,),
(self.batch_size,),
(self.num_entities,),
)
scores = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=self.num_entities // 2 + 1)
scores_no_slice = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=None)
self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice)
def _check_close_scores(self, scores, scores_no_slice):
self.assertTrue(torch.isfinite(scores).all(), msg=f"Normal scores had nan:\n\t{scores}")
self.assertTrue(torch.isfinite(scores_no_slice).all(), msg=f"Slice scores had nan\n\t{scores}")
self.assertTrue(torch.allclose(scores, scores_no_slice), msg=f"Differences: {scores - scores_no_slice}")
def _get_test_shapes(self) -> Collection[tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]]:
"""Return a set of test shapes for (h, r, t)."""
return (
( # single score
tuple(),
tuple(),
tuple(),
),
( # score_r with multi-t
(self.batch_size, 1, 1),
(1, self.num_relations, 1),
(self.batch_size, 1, self.num_entities // 2 + 1),
),
( # score_r with multi-t and broadcasted head
(1, 1, 1),
(1, self.num_relations, 1),
(self.batch_size, 1, self.num_entities),
),
( # full cwa
(self.num_entities, 1, 1),
(1, self.num_relations, 1),
(1, 1, self.num_entities),
),
)
def _get_output_shape(
self,
hs: tuple[int, ...],
rs: tuple[int, ...],
ts: tuple[int, ...],
) -> tuple[int, ...]:
components = []
if self.instance.entity_shape:
components.extend((hs, ts))
if self.instance.relation_shape:
components.append(rs)
return tuple(max(ds) for ds in zip(*components))
def test_forward(self):
"""Test forward."""
for hs, rs, ts in self._get_test_shapes():
if get_batchnorm_modules(self.instance) and any(numpy.prod(s) == 1 for s in (hs, rs, ts)):
logger.warning(
f"Skipping test for shapes {hs}, {rs}, {ts} because too small batch size for batch norm",
)
continue
h, r, t = self._get_hrt(hs, rs, ts)
scores = self.instance(h=h, r=r, t=t)
expected_shape = self._get_output_shape(hs, rs, ts)
self._check_scores(scores=scores, exp_shape=expected_shape)
def test_forward_consistency_with_functional(self):
"""Test forward's consistency with functional."""
if not isinstance(self.instance, FunctionalInteraction):
self.skipTest("Not a functional interaction")
# set in eval mode (otherwise there are non-deterministic factors like Dropout
self.instance.eval()
for hs, rs, ts in self._get_test_shapes():
h, r, t = self._get_hrt(hs, rs, ts)
scores = self.instance(h=h, r=r, t=t)
kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t)
scores_f = self.cls.func(**kwargs)
assert torch.allclose(scores, scores_f)
def test_scores(self):
"""Test individual scores."""
# set in eval mode (otherwise there are non-deterministic factors like Dropout
self.instance.eval()
for _ in range(10):
# test multiple different initializations
self.instance.reset_parameters()
h, r, t = self._get_hrt(tuple(), tuple(), tuple())
if isinstance(self.instance, FunctionalInteraction):
kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t)
# calculate by functional
scores_f = self.cls.func(**kwargs).view(-1)
else:
kwargs = dict(h=h, r=r, t=t)
scores_f = self.instance(h=h, r=r, t=t)
# calculate manually
scores_f_manual = self._exp_score(**kwargs).view(-1)
if not torch.allclose(scores_f, scores_f_manual, rtol=self.rtol, atol=self.atol):
# allclose checks: | input - other | < atol + rtol * |other|
a_delta = (scores_f_manual - scores_f).abs()
r_delta = (scores_f_manual - scores_f).abs() / scores_f.abs().clamp_min(1.0e-08)
raise AssertionError(
f"Abs. Diff: {a_delta.item()} (tol.: {self.atol}); Rel. Diff: {r_delta.item()} (tol. {self.rtol})",
)
@abstractmethod
def _exp_score(self, **kwargs) -> torch.FloatTensor:
"""Compute the expected score for a single-score batch."""
raise NotImplementedError(f"{self.cls.__name__}({sorted(kwargs.keys())})")
class TranslationalInteractionTests(InteractionTestCase, ABC):
"""Common tests for translational interaction."""
kwargs = dict(
p=2,
)
def _additional_score_checks(self, scores):
assert (scores <= 0).all()
class ResultTrackerTests(GenericTestCase[ResultTracker], unittest.TestCase):
"""Common tests for result trackers."""
def test_start_run(self):
"""Test start_run."""
self.instance.start_run(run_name="my_test.run")
def test_end_run(self):
"""Test end_run."""
self.instance.end_run()
def test_log_metrics(self):
"""Test log_metrics."""
for metrics, step, prefix in (
(
# simple
{"a": 1.0},
0,
None,
),
(
# nested
{"a": {"b": 5.0}, "c": -1.0},
2,
"test",
),
):
self.instance.log_metrics(metrics=metrics, step=step, prefix=prefix)
def test_log_params(self):
"""Test log_params."""
# nested
params = {
"num_epochs": 12,
"loss": {
"margin": 2.0, # a number
"normalize": True, # a bool
"activation": "relu", # a string
},
}
prefix = None
self.instance.log_params(params=params, prefix=prefix)
class FileResultTrackerTests(ResultTrackerTests):
"""Tests for FileResultTracker."""
def setUp(self) -> None:
"""Set up the file result tracker test."""
self.temporary_directory = tempfile.TemporaryDirectory()
self.path = pathlib.Path(self.temporary_directory.name).joinpath("test.log")
super().setUp()
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
# prepare a temporary test directory
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
kwargs["path"] = self.path
return kwargs
def tearDown(self) -> None: # noqa: D102
# check that file was created
assert self.path.is_file()
# make sure to close file before trying to delete it
self.instance.end_run()
# delete intermediate files
self.path.unlink()
self.temporary_directory.cleanup()
class RegularizerTestCase(GenericTestCase[Regularizer]):
"""A test case for quickly defining common tests for regularizers."""
#: The batch size
batch_size: int = 16
#: The device
device: torch.device
def post_instantiation_hook(self) -> None:
"""Move instance to device."""
self.device = resolve_device()
# move test instance to device
self.instance = self.instance.to(self.device)
def test_model(self) -> None:
"""Test whether the regularizer can be passed to a model."""
triples_factory = Nations().training
positive_batch = triples_factory.mapped_triples[: self.batch_size, :].to(device=self.device)
# Use RESCAL as it regularizes multiple tensors of different shape.
model = RESCAL(
triples_factory=triples_factory,
regularizer=self.instance,
).to(self.device)
# verify that the regularizer is stored for both, entity and relation representations
for r in (model.entity_representations, model.relation_representations):
assert len(r) == 1
self.assertEqual(r[0].regularizer, self.instance)
# Forward pass (should update regularizer)
model.score_hrt(hrt_batch=positive_batch)
# Call post_parameter_update (should reset regularizer)
model.post_parameter_update()
# Check if regularization term is reset
self.assertEqual(0.0, self.instance.term)
def _check_reset(self, instance: Optional[Regularizer] = None):
"""Verify that the regularizer is in resetted state."""
if instance is None:
instance = self.instance
# regularization term should be zero
self.assertEqual(0.0, instance.regularization_term.item())
# updated should be set to false
self.assertFalse(instance.updated)
def test_reset(self) -> None:
"""Test method `reset`."""
# call method
self.instance.reset()
self._check_reset()
def _generate_update_input(self, requires_grad: bool = False) -> Sequence[torch.FloatTensor]:
"""Generate input for update."""
# generate random tensors
return (
rand(self.batch_size, 10, generator=self.generator, device=self.device).requires_grad_(requires_grad),
rand(self.batch_size, 20, generator=self.generator, device=self.device).requires_grad_(requires_grad),
)
def _expected_updated_term(self, inputs: Sequence[torch.FloatTensor]) -> torch.FloatTensor:
"""Calculate the expected updated regularization term."""
exp_penalties = torch.stack([self._expected_penalty(x) for x in inputs])
expected_term = torch.sum(exp_penalties).view(1) * self.instance.weight
assert expected_term.shape == (1,)
return expected_term
def test_update(self) -> None:
"""Test method `update`."""
# generate inputs
inputs = self._generate_update_input()
# call update
self.instance.update(*inputs)
# check shape
self.assertEqual((1,), self.instance.term.shape)
# check result
expected_term = self._expected_updated_term(inputs=inputs)
self.assertAlmostEqual(self.instance.regularization_term.item(), expected_term.item())
def test_forward(self) -> None:
"""Test the regularizer's `forward` method."""
# generate single random tensor
x = rand(self.batch_size, 10, generator=self.generator, device=self.device)
# calculate penalty
penalty = self.instance(x=x)
# check shape
assert penalty.numel() == 1
# check value
expected_penalty = self._expected_penalty(x=x)
if expected_penalty is None:
logging.warning(f"{self.__class__.__name__} did not override `_expected_penalty`.")
else:
assert (expected_penalty == penalty).all()
def _expected_penalty(self, x: torch.FloatTensor) -> Optional[torch.FloatTensor]:
"""Compute expected penalty for given tensor."""
return None
def test_pop_regularization_term(self):
"""Verify popping a regularization term."""
# update term
inputs = self._generate_update_input(requires_grad=True)
self.instance.update(*inputs)
# check that the expected term is returned
exp = (self.instance.weight * self._expected_updated_term(inputs)).item()
self.assertEqual(exp, self.instance.pop_regularization_term().item())
# check that the regularizer is now reset
self._check_reset()
def test_apply_only_once(self):
"""Test apply-only-once support."""
# create another instance with apply_only_once enabled
instance = self.cls(**ChainMap(dict(apply_only_once=True), self.instance_kwargs)).to(self.device)
# test initial state
self._check_reset(instance=instance)
# after first update, should change the term
first_tensors = self._generate_update_input()
instance.update(*first_tensors)
self.assertTrue(instance.updated)
self.assertNotEqual(0.0, instance.regularization_term.item())
term = instance.regularization_term.clone()
# after second update, no change should happen
second_tensors = self._generate_update_input()
instance.update(*second_tensors)
self.assertTrue(instance.updated)
self.assertEqual(term, instance.regularization_term)
class LpRegularizerTest(RegularizerTestCase):
"""Common test for L_p regularizers."""
cls = LpRegularizer
def _expected_penalty(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102
kwargs = self.kwargs
if kwargs is None:
kwargs = {}
p = kwargs.get("p", self.instance.p)
value = x.norm(p=p, dim=-1).mean()
if kwargs.get("normalize", False):
dim = torch.as_tensor(x.shape[-1], dtype=torch.float, device=x.device)
# FIXME isn't any finite number allowed now?
if p == 2:
value = value / dim.sqrt()
elif p == 1:
value = value / dim
else:
raise NotImplementedError
return value
class ModelTestCase(unittest_templates.GenericTestCase[Model]):
"""A test case for quickly defining common tests for KGE models."""
#: Additional arguments passed to the training loop's constructor method
training_loop_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None
#: The triples factory instance
factory: TriplesFactory
#: The batch size for use for forward_* tests
batch_size: int = 20
#: The embedding dimensionality
embedding_dim: int = 3
#: Whether to create inverse triples (needed e.g. by ConvE)
create_inverse_triples: bool = False
#: The sampler to use for sLCWA (different e.g. for R-GCN)
sampler: Optional[str] = None
#: The batch size for use when testing training procedures
train_batch_size = 400
#: The number of epochs to train the model
train_num_epochs = 2
#: A random number generator from torch
generator: torch.Generator
#: The number of parameters which receive a constant (i.e. non-randomized)
# initialization
num_constant_init: int = 0
#: Static extras to append to the CLI
cli_extras: Sequence[str] = tuple()
#: the model's device
device: torch.device
#: the inductive mode
mode: ClassVar[Optional[InductiveMode]] = None
def pre_setup_hook(self) -> None: # noqa: D102
# for reproducible testing
_, self.generator, _ = set_random_seed(42)
self.device = resolve_device()
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
dataset = Nations(create_inverse_triples=self.create_inverse_triples)
self.factory = dataset.training
# insert shared parameters
kwargs["triples_factory"] = self.factory
kwargs["embedding_dim"] = self.embedding_dim
return kwargs
def post_instantiation_hook(self) -> None: # noqa: D102
# move model to correct device
self.instance = self.instance.to(self.device)
def test_get_grad_parameters(self):
"""Test the model's ``get_grad_params()`` method."""
self.assertLess(
0, len(list(self.instance.get_grad_params())), msg="There is not at least one trainable parameter"
)
# Check that all the parameters actually require a gradient
for parameter in self.instance.get_grad_params():
assert parameter.requires_grad
# Try to initialize an optimizer
optimizer = SGD(params=self.instance.get_grad_params(), lr=1.0)
assert optimizer is not None
def test_reset_parameters_(self):
"""Test :func:`Model.reset_parameters_`."""
# get model parameters
params = list(self.instance.parameters())
old_content = {id(p): p.data.detach().clone() for p in params}
# re-initialize
self.instance.reset_parameters_()
# check that the operation works in-place
new_params = list(self.instance.parameters())
assert set(id(np) for np in new_params) == set(id(p) for p in params)
# check that the parameters where modified
num_equal_weights_after_re_init = sum(
1 for new_param in new_params if (new_param.data == old_content[id(new_param)]).all()
)
self.assertEqual(num_equal_weights_after_re_init, self.num_constant_init)
def _check_scores(self, batch, scores) -> None:
"""Check the scores produced by a forward function."""
# check for finite values by default
self.assertTrue(torch.all(torch.isfinite(scores)).item(), f"Some scores were not finite:\n{scores}")
# check whether a gradient can be back-propgated
scores.mean().backward()
def test_save(self) -> None:
"""Test that the model can be saved properly."""
with tempfile.TemporaryDirectory() as temp_directory:
torch.save(self.instance, os.path.join(temp_directory, "model.pickle"))
def _test_score(
self, score: Callable, columns: Union[Sequence[int], slice], shape: tuple[int, ...], **kwargs
) -> None:
"""Test score functions."""
batch = self.factory.mapped_triples[: self.batch_size, columns].to(self.instance.device)
try:
scores = score(batch, mode=self.mode, **kwargs)
except ValueError as error:
raise SkipTest() from error
except NotImplementedError:
self.fail(msg=f"{score} not yet implemented")
except RuntimeError as e:
if str(e) == "fft: ATen not compiled with MKL support":
self.skipTest(str(e))
else:
raise e
if score is self.instance.score_r and self.create_inverse_triples:
# TODO: look into score_r for inverse relations
logger.warning("score_r's shape is not clear yet for models with inverse relations")
else:
self.assertTupleEqual(tuple(scores.shape), shape)
self._check_scores(batch, scores)
# clear buffers for message passing models
self.instance.post_parameter_update()
def _test_score_multi(self, name: str, max_id: int, **kwargs):
"""Test score functions with multi scoring."""
k = max_id // 2
for ids in (
torch.randperm(max_id)[:k],
torch.randint(max_id, size=(self.batch_size, k)),
):
with self.subTest(shape=ids.shape):
self._test_score(shape=(self.batch_size, k), **kwargs, **{name: ids.to(device=self.instance.device)})
def test_score_hrt(self) -> None:
"""Test the model's ``score_hrt()`` function."""
self._test_score(score=self.instance.score_hrt, columns=slice(None), shape=(self.batch_size, 1))
def test_score_t(self) -> None:
"""Test the model's ``score_t()`` function."""
self._test_score(
score=self.instance.score_t, columns=slice(0, 2), shape=(self.batch_size, self.instance.num_entities)
)
def test_score_t_multi(self) -> None:
"""Test the model's ``score_t()`` function with custom tail candidates."""
self._test_score_multi(
name="tails", max_id=self.factory.num_entities, score=self.instance.score_t, columns=slice(0, 2)
)
def test_score_r(self) -> None:
"""Test the model's ``score_r()`` function."""
self._test_score(
score=self.instance.score_r,
columns=[0, 2],
shape=(self.batch_size, self.instance.num_relations),
)
def test_score_r_multi(self) -> None:
"""Test the model's ``score_r()`` function with custom relation candidates."""
self._test_score_multi(
name="relations", max_id=self.factory.num_relations, score=self.instance.score_r, columns=[0, 2]
)
def test_score_h(self) -> None:
"""Test the model's ``score_h()`` function."""
self._test_score(
score=self.instance.score_h, columns=slice(1, None), shape=(self.batch_size, self.instance.num_entities)
)
def test_score_h_multi(self) -> None:
"""Test the model's ``score_h()`` function with custom head candidates."""
self._test_score_multi(
name="heads", max_id=self.factory.num_entities, score=self.instance.score_h, columns=slice(1, None)
)
@pytest.mark.slow
def test_train_slcwa(self) -> None:
"""Test that sLCWA training does not fail."""
loop = SLCWATrainingLoop(
model=self.instance,
triples_factory=self.factory,
optimizer=Adagrad(params=self.instance.get_grad_params(), lr=0.001),
**(self.training_loop_kwargs or {}),
)
losses = self._safe_train_loop(
loop,
num_epochs=self.train_num_epochs,
batch_size=self.train_batch_size,
sampler=self.sampler,
)
self.assertIsInstance(losses, list)
@pytest.mark.slow
def test_train_lcwa(self) -> None:
"""Test that LCWA training does not fail."""
loop = LCWATrainingLoop(
model=self.instance,
triples_factory=self.factory,
optimizer=Adagrad(params=self.instance.get_grad_params(), lr=0.001),
**(self.training_loop_kwargs or {}),
)
losses = self._safe_train_loop(
loop,
num_epochs=self.train_num_epochs,
batch_size=self.train_batch_size,
sampler=None,
)
self.assertIsInstance(losses, list)
def _safe_train_loop(self, loop: TrainingLoop, num_epochs, batch_size, sampler):
try:
losses = loop.train(
triples_factory=self.factory,
num_epochs=num_epochs,
batch_size=batch_size,
sampler=sampler,
use_tqdm=False,
)
except RuntimeError as e:
if str(e) == "fft: ATen not compiled with MKL support":
self.skipTest(str(e))
else:
raise e
else:
return losses
def test_save_load_model_state(self):
"""Test whether a saved model state can be re-loaded."""
original_model = self.cls(
random_seed=42,
**self.instance_kwargs,
)
loaded_model = self.cls(
random_seed=21,
**self.instance_kwargs,
)
def _equal_embeddings(a: Representation, b: Representation) -> bool:
"""Test whether two embeddings are equal."""
return (a(indices=None) == b(indices=None)).all()
with tempfile.TemporaryDirectory() as tmpdirname:
file_path = os.path.join(tmpdirname, "test.pt")
original_model.save_state(path=file_path)
loaded_model.load_state(path=file_path)
@property
def _cli_extras(self):
"""Return a list of extra flags for the CLI."""
kwargs = self.kwargs or {}
extras = [
"--silent",
]
for k, v in kwargs.items():
extras.append("--" + k.replace("_", "-"))
extras.append(str(v))
# For the high/low memory test cases of NTN, SE, etc.
if self.training_loop_kwargs and "automatic_memory_optimization" in self.training_loop_kwargs:
automatic_memory_optimization = self.training_loop_kwargs.get("automatic_memory_optimization")
if automatic_memory_optimization is True:
extras.append("--automatic-memory-optimization")
elif automatic_memory_optimization is False:
extras.append("--no-automatic-memory-optimization")
# else, leave to default
extras += [
"--number-epochs",
self.train_num_epochs,
"--embedding-dim",
self.embedding_dim,
"--batch-size",
self.train_batch_size,
]
extras.extend(self.cli_extras)
# Make sure that inverse triples are created if create_inverse_triples=True
if self.create_inverse_triples:
extras.append("--create-inverse-triples")
extras = [str(e) for e in extras]
return extras
@pytest.mark.slow
def test_cli_training_nations(self):
"""Test running the pipeline on almost all models with only training data."""
self._help_test_cli(["-t", NATIONS_TRAIN_PATH] + self._cli_extras)
@pytest.mark.slow
def test_pipeline_nations_early_stopper(self):
"""Test running the pipeline with early stopping."""
model_kwargs = dict(self.instance_kwargs)
# triples factory is added by the pipeline
model_kwargs.pop("triples_factory")
pipeline(
model=self.cls,
model_kwargs=model_kwargs,
dataset="nations",
dataset_kwargs=dict(create_inverse_triples=self.create_inverse_triples),
stopper="early",
training_loop_kwargs=self.training_loop_kwargs,
stopper_kwargs=dict(frequency=1),
training_kwargs=dict(
batch_size=self.train_batch_size,
num_epochs=self.train_num_epochs,
),
)
@pytest.mark.slow
def test_cli_training_kinships(self):
"""Test running the pipeline on almost all models with only training data."""
self._help_test_cli(["-t", KINSHIPS_TRAIN_PATH] + self._cli_extras)
@pytest.mark.slow
def test_cli_training_nations_testing(self):
"""Test running the pipeline on almost all models with only training data."""
self._help_test_cli(["-t", NATIONS_TRAIN_PATH, "-q", NATIONS_TEST_PATH] + self._cli_extras)
def _help_test_cli(self, args):
"""Test running the pipeline on all models."""
if (
issubclass(self.cls, (pykeen.models.RGCN, pykeen.models.CooccurrenceFilteredModel))
or self.cls is pykeen.models.ERModel
):
self.skipTest(f"Cannot choose interaction via CLI for {self.cls}.")
runner = CliRunner()
cli = build_cli_from_cls(self.cls)
# TODO: Catch HolE MKL error?
result: Result = runner.invoke(cli, args)
self.assertEqual(
0,
result.exit_code,
msg=f"""
Command
=======
$ pykeen train {self.cls.__name__.lower()} {' '.join(map(str, args))}
Output
======
{result.output}
Exception
=========
{result.exc_info[1]}
Traceback
=========
{''.join(traceback.format_tb(result.exc_info[2]))}
""",
)
def test_has_hpo_defaults(self):
"""Test that there are defaults for HPO."""
try:
d = self.cls.hpo_default
except AttributeError:
self.fail(msg=f"{self.cls.__name__} is missing hpo_default class attribute")
else:
self.assertIsInstance(d, dict)
def test_post_parameter_update_regularizer(self):
"""Test whether post_parameter_update resets the regularization term."""
if not hasattr(self.instance, "regularizer"):
self.skipTest("no regularizer")
# set regularizer term to something that isn't zero
self.instance.regularizer.regularization_term = torch.ones(1, dtype=torch.float, device=self.instance.device)
# call post_parameter_update
self.instance.post_parameter_update()
# assert that the regularization term has been reset
expected_term = torch.zeros(1, dtype=torch.float, device=self.instance.device)
assert self.instance.regularizer.regularization_term == expected_term
def test_post_parameter_update(self):
"""Test whether post_parameter_update correctly enforces model constraints."""
# do one optimization step
opt = optim.SGD(params=self.instance.parameters(), lr=1.0)
batch = self.factory.mapped_triples[: self.batch_size, :].to(self.instance.device)
scores = self.instance.score_hrt(hrt_batch=batch, mode=self.mode)
fake_loss = scores.mean()
fake_loss.backward()
opt.step()
# call post_parameter_update
self.instance.post_parameter_update()
# check model constraints
self._check_constraints()
def _check_constraints(self):
"""Check model constraints."""
def _test_score_equality(self, columns: Union[slice, list[int]], name: str) -> None:
"""Migration tests for non-ERModel models testing for consistent optimized score implementations."""
if isinstance(self.instance, ERModel):
raise SkipTest("ERModel fulfils this by design.")
if isinstance(self.instance, CooccurrenceFilteredModel):
raise SkipTest("CooccurrenceFilteredModel fulfils this if its base model fulfils it.")
batch = self.factory.mapped_triples[: self.batch_size, columns].to(self.instance.device)
self.instance.eval()
try:
scores = getattr(self.instance, name)(batch)
scores_super = getattr(super(self.instance.__class__, self.instance), name)(batch)
except NotImplementedError:
self.fail(msg=f"{name} not yet implemented")
except RuntimeError as e:
if str(e) == "fft: ATen not compiled with MKL support":
self.skipTest(str(e))
else:
raise e
self.assertIsNotNone(scores)
self.assertIsNotNone(scores_super)
assert torch.allclose(scores, scores_super, atol=1e-06)
def test_score_h_with_score_hrt_equality(self) -> None:
"""Test the equality of the model's ``score_h()`` and ``score_hrt()`` function."""
self._test_score_equality(columns=slice(1, None), name="score_h")
def test_score_r_with_score_hrt_equality(self) -> None:
"""Test the equality of the model's ``score_r()`` and ``score_hrt()`` function."""
self._test_score_equality(columns=[0, 2], name="score_r")
def test_score_t_with_score_hrt_equality(self) -> None:
"""Test the equality of the model's ``score_t()`` and ``score_hrt()`` function."""
self._test_score_equality(columns=slice(2), name="score_t")
def test_reset_parameters_constructor_call(self):
"""Tests whether reset_parameters is called in the constructor."""
with patch.object(self.cls, "reset_parameters_", return_value=None) as mock_method:
try:
self.cls(**self.instance_kwargs)
except TypeError as error:
assert error.args == ("'NoneType' object is not callable",)
mock_method.assert_called_once()
class DistanceModelTestCase(ModelTestCase):
"""A test case for distance-based models."""
def _check_scores(self, batch, scores) -> None:
super()._check_scores(batch=batch, scores=scores)
# Distance-based model
assert (scores <= 0.0).all()
class BaseKG2ETest(ModelTestCase):
"""General tests for the KG2E model."""
cls = pykeen.models.KG2E
c_min: float = 0.01
c_max: float = 1.0
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
kwargs["c_min"] = self.c_min
kwargs["c_max"] = self.c_max
return kwargs
def _check_constraints(self):
"""Check model constraints.
* Entity and relation embeddings have to have at most unit L2 norm.
* Covariances have to have values between c_min and c_max
"""
self.instance: ERModel
(e_mean, e_cov), (r_mean, r_cov) = self.instance.entity_representations, self.instance.relation_representations
for embedding in (e_mean, r_mean):
assert all_in_bounds(embedding(indices=None).norm(p=2, dim=-1), high=1.0, a_tol=EPSILON)
for cov in (e_cov, r_cov):
assert all_in_bounds(
cov(indices=None), low=self.instance_kwargs["c_min"], high=self.instance_kwargs["c_max"]
)
class BaseRGCNTest(ModelTestCase):
"""Test the R-GCN model."""
cls = pykeen.models.RGCN
sampler = "schlichtkrull"
def _check_constraints(self):
"""Check model constraints.
Enriched embeddings have to be reset.
"""
assert self.instance.entity_representations[0].enriched_embeddings is None
class BaseNodePieceTest(ModelTestCase):
"""Test the NodePiece model."""
cls = pykeen.models.NodePiece
create_inverse_triples = True
def _help_test_cli(self, args): # noqa: D102
if self.instance_kwargs.get("tokenizers_kwargs"):
raise SkipTest("No support for tokenizers_kwargs via CLI.")
return super()._help_test_cli(args)
class InductiveModelTestCase(ModelTestCase):
"""Tests for inductive models."""
mode = TRAINING
num_relations: ClassVar[int] = 7
num_entities_transductive: ClassVar[int] = 13
num_entities_inductive: ClassVar[int] = 5
num_triples_training: ClassVar[int] = 33
num_triples_inference: ClassVar[int] = 31
num_triples_testing: ClassVar[int] = 37
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
dataset = create_inductive_dataset(
num_relations=self.num_relations,
num_entities_transductive=self.num_entities_transductive,
num_entities_inductive=self.num_entities_inductive,
num_triples_training=self.num_triples_training,
num_triples_inference=self.num_triples_inference,
num_triples_testing=self.num_triples_testing,
create_inverse_triples=self.create_inverse_triples,
)
training_loop_kwargs = dict(self.training_loop_kwargs or dict())
training_loop_kwargs["mode"] = self.mode
InductiveModelTestCase.training_loop_kwargs = training_loop_kwargs
# dataset = InductiveFB15k237(create_inverse_triples=self.create_inverse_triples)
kwargs["triples_factory"] = self.factory = dataset.transductive_training
kwargs["inference_factory"] = dataset.inductive_inference
return kwargs
def _help_test_cli(self, args): # noqa: D102
raise SkipTest("Inductive models are not compatible the CLI.")
def test_pipeline_nations_early_stopper(self): # noqa: D102
raise SkipTest("Inductive models are not compatible the pipeline.")
class RepresentationTestCase(GenericTestCase[Representation]):
"""Common tests for representation modules."""
batch_size: ClassVar[int] = 2
num_negatives: ClassVar[int] = 3
max_id: ClassVar[int] = 7
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
kwargs = super()._pre_instantiation_hook(kwargs)
kwargs.update(dict(max_id=self.max_id))
return kwargs
def _check_result(self, x: torch.FloatTensor, prefix_shape: tuple[int, ...]):
"""Check the result."""
# check type
assert torch.is_tensor(x)
assert x.dtype == torch.get_default_dtype()
# check shape
expected_shape = prefix_shape + self.instance.shape
self.assertEqual(x.shape, expected_shape)
def _test_forward(self, indices: Optional[torch.LongTensor]):
"""Test forward method."""
representations = self.instance(indices=indices)
prefix_shape = (self.instance.max_id,) if indices is None else tuple(indices.shape)
self._check_result(x=representations, prefix_shape=prefix_shape)
def _test_indices(self, indices: Optional[torch.LongTensor]):
"""Test forward and canonical shape for indices."""
self._test_forward(indices=indices)
def test_max_id(self):
"""Test maximum id."""
self.assertEqual(self.max_id, self.instance.max_id)
def test_no_indices(self):
"""Test without indices."""
self._test_indices(indices=None)
def test_1d_indices(self):
"""Test with 1-dimensional indices."""
self._test_indices(indices=torch.randint(self.instance.max_id, size=(self.batch_size,)))
def test_2d_indices(self):
"""Test with 1-dimensional indices."""
self._test_indices(indices=(torch.randint(self.instance.max_id, size=(self.batch_size, self.num_negatives))))
def test_all_indices(self):
"""Test with all indices."""
self._test_indices(indices=torch.arange(self.instance.max_id))
def test_dropout(self):
"""Test dropout layer."""
# create a new instance with guaranteed dropout
kwargs = self.instance_kwargs
kwargs.pop("dropout", None)
dropout_instance = self.cls(**kwargs, dropout=0.1)
# set to training mode
dropout_instance.train()
# check for different output
indices = torch.arange(2)
# use more samples to make sure that enough values can be dropped
a = torch.stack([dropout_instance(indices) for _ in range(20)])
assert not (a[0:1] == a).all()
def test_str(self):
"""Test generating the string representation."""
# this implicitly tests extra_repr / iter_extra_repr
assert isinstance(str(self), str)
class TriplesFactoryRepresentationTestCase(RepresentationTestCase):
"""Tests for representations requiring triples factories."""
num_entities: ClassVar[int]
num_relations: ClassVar[int] = 7
num_triples: ClassVar[int] = 31
create_inverse_triples: bool = False
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
self.num_entities = self.max_id
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
kwargs["triples_factory"] = generation.generate_triples_factory(
num_entities=self.max_id,
num_relations=self.num_relations,
num_triples=self.num_triples,
create_inverse_triples=self.create_inverse_triples,
)
return kwargs
@needs_packages("torch_geometric")
class MessagePassingRepresentationTests(TriplesFactoryRepresentationTestCase):
"""Tests for message passing representations."""
def test_consistency_k_hop(self):
"""Test consistency of results between using only k-hop and using the full graph."""
# select random indices
indices = torch.randint(self.num_entities, size=(self.num_entities // 2,), generator=self.generator)
assert isinstance(self.instance, pykeen.nn.pyg.MessagePassingRepresentation)
# forward pass with full graph
self.instance.restrict_k_hop = False
x_full = self.instance(indices=indices)
# forward pass with restricted graph
self.instance.restrict_k_hop = True
x_restrict = self.instance(indices=indices)
# verify the results are similar
assert torch.allclose(x_full, x_restrict)
class EdgeWeightingTestCase(GenericTestCase[pykeen.nn.weighting.EdgeWeighting]):
"""Tests for message weighting."""
#: The number of entities
num_entities: int = 16
#: The number of triples
num_triples: int = 101
#: the message dim
message_dim: int = 3
def post_instantiation_hook(self): # noqa: D102
self.source, self.target = torch.randint(self.num_entities, size=(2, self.num_triples))
self.message = torch.rand(self.num_triples, self.message_dim, requires_grad=True)
# TODO: separation message vs. entity dim?
self.x_e = torch.rand(self.num_entities, self.message_dim)
def _test(self, weights: torch.FloatTensor, shape: tuple[int, ...]):
"""Perform common tests."""
# check shape
assert weights.shape == shape
# check dtype
assert weights.dtype == torch.float32
# check finite values (e.g. due to division by zero)
assert torch.isfinite(weights).all()
# check non-negativity
assert (weights >= 0.0).all()
def test_message_weighting(self):
"""Test message weighting with message."""
self._test(
weights=self.instance(source=self.source, target=self.target, message=self.message, x_e=self.x_e),
shape=self.message.shape,
)
def test_message_weighting_no_message(self):
"""Test message weighting without message."""
if self.instance.needs_message:
raise SkipTest(f"{self.cls} needs messages for weighting them.")
self._test(weights=self.instance(source=self.source, target=self.target), shape=self.source.shape)
class DecompositionTestCase(GenericTestCase[pykeen.nn.message_passing.Decomposition]):
"""Tests for relation-specific weight decomposition message passing classes."""
#: the input dimension
input_dim: int = 8
#: the output dimension
output_dim: int = 4
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
self.factory = Nations().training
self.source, self.edge_type, self.target = self.factory.mapped_triples.t()
self.x = torch.rand(self.factory.num_entities, self.input_dim, requires_grad=True)
kwargs["input_dim"] = self.input_dim
kwargs["output_dim"] = self.output_dim
kwargs["num_relations"] = self.factory.num_relations
return kwargs
def test_forward(self):
"""Test the :meth:`Decomposition.forward` function."""
for edge_weights in [None, torch.rand_like(self.source, dtype=torch.get_default_dtype())]:
y = self.instance(
x=self.x,
source=self.source,
target=self.target,
edge_type=self.edge_type,
edge_weights=edge_weights,
)
assert y.shape == (self.x.shape[0], self.output_dim)
def prepare_adjacency(self, horizontal: bool) -> torch.Tensor:
"""
Prepare adjacency matrix for the given stacking direction.
:param horizontal:
whether to stack horizontally or vertically
:return:
the adjacency matrix
"""
return adjacency_tensor_to_stacked_matrix(
num_relations=self.factory.num_relations,
num_entities=self.factory.num_entities,
source=self.source,
target=self.target,
edge_type=self.edge_type,
horizontal=horizontal,
)
def check_output(self, x: torch.Tensor):
"""Check the output tensor."""
assert torch.is_tensor(x)
assert x.shape == (self.factory.num_entities, self.output_dim)
assert x.requires_grad
def test_horizontal(self):
"""Test processing of horizontally stacked matrix."""
adj = self.prepare_adjacency(horizontal=True)
x = self.instance.forward_horizontally_stacked(x=self.x, adj=adj)
self.check_output(x=x)
def test_vertical(self):
"""Test processing of vertically stacked matrix."""
adj = self.prepare_adjacency(horizontal=False)
x = self.instance.forward_vertically_stacked(x=self.x, adj=adj)
self.check_output(x=x)
class InitializerTestCase(unittest.TestCase):
"""A test case for initializers."""
#: the number of entities
num_entities: ClassVar[int] = 33
#: the shape of the tensor to initialize
shape: ClassVar[tuple[int, ...]] = (3,)
#: to be initialized / set in subclass
initializer: Initializer
#: the interaction to use for testing a model
interaction: ClassVar[HintOrType[Interaction]] = DistMultInteraction
dtype: ClassVar[torch.dtype] = torch.get_default_dtype()
def test_initialization(self):
"""Test whether the initializer returns a modified tensor."""
shape = (self.num_entities, *self.shape)
if self.dtype.is_complex:
shape = shape + (2,)
x = torch.rand(size=shape)
# initializers *may* work in-place => clone
y = self.initializer(x.clone())
assert not (x == y).all()
self._verify_initialization(y)
def _verify_initialization(self, x: torch.FloatTensor) -> None:
"""Verify properties of initialization."""
pass
def test_model(self):
"""Test whether initializer can be used for a model."""
triples_factory = generation.generate_triples_factory(num_entities=self.num_entities)
# actual number may be different...
self.num_entities = triples_factory.num_entities
model = pykeen.models.ERModel(
triples_factory=triples_factory,
interaction=self.interaction,
entity_representations_kwargs=dict(shape=self.shape, initializer=self.initializer, dtype=self.dtype),
relation_representations_kwargs=dict(shape=self.shape),
random_seed=0,
).to(resolve_device())
model.reset_parameters_()
class PredictBaseTestCase(unittest.TestCase):
"""Base test for prediction workflows."""
batch_size: ClassVar[int] = 2
model_cls: ClassVar[type[Model]]
model_kwargs: ClassVar[Mapping[str, Any]]
factory: TriplesFactory
batch: MappedTriples
model: Model
def setUp(self) -> None:
"""Prepare model."""
self.factory = Nations().training
self.batch = self.factory.mapped_triples[: self.batch_size, :]
self.model = self.model_cls(
triples_factory=self.factory,
**self.model_kwargs,
)
class CleanerTestCase(GenericTestCase[Cleaner]):
"""Test cases for cleaner."""
def post_instantiation_hook(self) -> None:
"""Prepare triples."""
self.dataset = Nations()
self.all_entities = set(range(self.dataset.num_entities))
self.mapped_triples = self.dataset.training.mapped_triples
# unfavourable split to ensure that cleanup is necessary
self.reference, self.other = torch.split(
self.mapped_triples,
split_size_or_sections=[24, self.mapped_triples.shape[0] - 24],
dim=0,
)
# check for unclean split
assert get_entities(self.reference) != self.all_entities
def test_cleanup_pair(self):
"""Test cleanup_pair."""
reference_clean, other_clean = self.instance.cleanup_pair(
reference=self.reference,
other=self.other,
random_state=42,
)
# check that no triple got lost
assert triple_tensor_to_set(self.mapped_triples) == triple_tensor_to_set(
torch.cat(
[
reference_clean,
other_clean,
],
dim=0,
)
)
# check that triples where only moved from other to reference
assert is_triple_tensor_subset(self.reference, reference_clean)
assert is_triple_tensor_subset(other_clean, self.other)
# check that all entities occur in reference
assert get_entities(reference_clean) == self.all_entities
def test_call(self):
"""Test call."""
triples_groups = [self.reference] + list(torch.split(self.other, split_size_or_sections=3, dim=0))
clean_groups = self.instance(triples_groups=triples_groups, random_state=42)
assert all(torch.is_tensor(triples) and triples.dtype for triples in clean_groups)
class SplitterTestCase(GenericTestCase[Splitter]):
"""Test cases for triples splitter."""
def post_instantiation_hook(self) -> None:
"""Prepare data."""
dataset = Nations()
self.all_entities = set(range(dataset.num_entities))
self.mapped_triples = dataset.training.mapped_triples
def _test_split(self, ratios: Union[float, Sequence[float]], exp_parts: int):
"""Test splitting."""
splitted = self.instance.split(
mapped_triples=self.mapped_triples,
ratios=ratios,
random_state=None,
)
assert len(splitted) == exp_parts
# check that no triple got lost
assert triple_tensor_to_set(self.mapped_triples) == set().union(
*(triple_tensor_to_set(triples) for triples in splitted)
)
# check that all entities are covered in first part
assert triple_tensor_to_set(splitted[0]) == self.all_entities
class EvaluatorTestCase(unittest_templates.GenericTestCase[Evaluator]):
"""A test case for quickly defining common tests for evaluators models."""
# the model
model: Model
# Settings
batch_size: int = 8
embedding_dim: int = 7
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
self.dataset = Nations()
return super()._pre_instantiation_hook(kwargs=kwargs)
@property
def factory(self) -> CoreTriplesFactory:
"""Return the evaluation factory."""
return self.dataset.validation
def post_instantiation_hook(self) -> None: # noqa: D102
# Use small model (untrained)
self.model = TransE(triples_factory=self.factory, embedding_dim=self.embedding_dim)
def _get_input(
self,
inverse: bool = False,
) -> tuple[torch.LongTensor, torch.FloatTensor, Optional[torch.BoolTensor]]:
# Get batch
hrt_batch = self.factory.mapped_triples[: self.batch_size].to(self.model.device)
# Compute scores
if inverse:
scores = self.model.score_h(rt_batch=hrt_batch[:, 1:])
else:
scores = self.model.score_t(hr_batch=hrt_batch[:, :2])
# Compute mask only if required
if self.instance.requires_positive_mask:
# TODO: Re-use filtering code
triples = self.factory.mapped_triples.to(self.model.device)
if inverse:
sel_col, start_col = 0, 1
else:
sel_col, start_col = 2, 0
stop_col = start_col + 2
# shape: (batch_size, num_triples)
triple_mask = (triples[None, :, start_col:stop_col] == hrt_batch[:, None, start_col:stop_col]).all(dim=-1)
batch_indices, triple_indices = triple_mask.nonzero(as_tuple=True)
entity_indices = triples[triple_indices, sel_col]
# shape: (batch_size, num_entities)
mask = torch.zeros_like(scores, dtype=torch.bool)
mask[batch_indices, entity_indices] = True
else:
mask = None
return hrt_batch, scores, mask
def test_process_tail_scores_(self) -> None:
"""Test the evaluator's ``process_tail_scores_()`` function."""
hrt_batch, scores, mask = self._get_input()
true_scores = scores[torch.arange(0, hrt_batch.shape[0]), hrt_batch[:, 2]][:, None]
self.instance.process_scores_(
hrt_batch=hrt_batch,
target=LABEL_TAIL,
true_scores=true_scores,
scores=scores,
dense_positive_mask=mask,
)
def test_process_head_scores_(self) -> None:
"""Test the evaluator's ``process_head_scores_()`` function."""
hrt_batch, scores, mask = self._get_input(inverse=True)
true_scores = scores[torch.arange(0, hrt_batch.shape[0]), hrt_batch[:, 0]][:, None]
self.instance.process_scores_(
hrt_batch=hrt_batch,
target=LABEL_HEAD,
true_scores=true_scores,
scores=scores,
dense_positive_mask=mask,
)
def _process_batches(self):
"""Process one batch per side."""
hrt_batch, scores, mask = self._get_input()
true_scores = scores[torch.arange(0, hrt_batch.shape[0]), hrt_batch[:, 2]][:, None]
for target in (LABEL_HEAD, LABEL_TAIL):
self.instance.process_scores_(
hrt_batch=hrt_batch,
target=target,
true_scores=true_scores,
scores=scores,
dense_positive_mask=mask,
)
return hrt_batch, scores, mask
def test_finalize(self) -> None:
"""Test the finalize() function."""
# Process one batch
hrt_batch, scores, mask = self._process_batches()
result = self.instance.finalize()
assert isinstance(result, MetricResults)
self._validate_result(
result=result,
data={"batch": hrt_batch, "scores": scores, "mask": mask},
)
def _validate_result(
self,
result: MetricResults,
data: dict[str, torch.Tensor],
):
logger.warning(f"{self.__class__.__name__} did not overwrite _validate_result.")
def test_pipeline(self):
"""Test interaction with pipeline."""
pipeline(
training=self.factory,
testing=self.factory,
model="distmult",
evaluator=evaluator_resolver.normalize_cls(self.cls),
evaluator_kwargs=self.instance_kwargs,
training_kwargs=dict(
num_epochs=1,
),
)
class AnchorSelectionTestCase(GenericTestCase[pykeen.nn.node_piece.AnchorSelection]):
"""Tests for anchor selection."""
num_anchors: int = 7
num_entities: int = 33
num_triples: int = 101
edge_index: numpy.ndarray
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
"""Prepare kwargs."""
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
kwargs["num_anchors"] = self.num_anchors
return kwargs
def post_instantiation_hook(self) -> None:
"""Prepare edge index."""
generator = numpy.random.default_rng(seed=42)
self.edge_index = generator.integers(low=0, high=self.num_entities, size=(2, self.num_triples))
def test_call(self):
"""Test __call__."""
anchors = self.instance(edge_index=self.edge_index)
# shape
assert len(anchors) == self.num_anchors
# value range
assert (0 <= anchors).all()
assert (anchors < self.num_entities).all()
# no duplicates
assert len(set(anchors.tolist())) == len(anchors)
class AnchorSearcherTestCase(GenericTestCase[pykeen.nn.node_piece.AnchorSearcher]):
"""Tests for anchor search."""
num_entities = 33
k: int = 2
edge_index: numpy.ndarray
anchors: numpy.ndarray
def post_instantiation_hook(self) -> None:
"""Prepare circular edge index."""
self.edge_index = numpy.stack(
[numpy.arange(self.num_entities), (numpy.arange(self.num_entities) + 1) % self.num_entities],
axis=0,
)
self.anchors = numpy.arange(0, self.num_entities, 10)
def test_call(self):
"""Test __call__."""
tokens = self.instance(edge_index=self.edge_index, anchors=self.anchors, k=self.k)
# shape
assert tokens.shape == (self.num_entities, self.k)
# value range
assert (tokens >= -1).all()
assert (tokens < len(self.anchors)).all()
# no duplicates
for row in tokens.tolist():
self.assertDictEqual({k: v for k, v in Counter(row).items() if k >= 0 and v > 1}, {}, msg="duplicate token")
class TokenizerTestCase(GenericTestCase[pykeen.nn.node_piece.Tokenizer]):
"""Tests for tokenization."""
num_tokens: int = 2
factory: CoreTriplesFactory
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
"""Prepare triples."""
self.factory = Nations().training
return {}
def test_call(self):
"""Test __call__."""
vocabulary_size, tokens = self.instance(
mapped_triples=self.factory.mapped_triples,
num_tokens=self.num_tokens,
num_entities=self.factory.num_entities,
num_relations=self.factory.num_relations,
)
assert isinstance(vocabulary_size, int)
assert vocabulary_size > 0
# shape
assert tokens.shape == (self.factory.num_entities, self.num_tokens)
# value range
assert (tokens >= -1).all()
# no repetition, except padding idx
for row in tokens.tolist():
self.assertDictEqual({k: v for k, v in Counter(row).items() if k >= 0 and v > 1}, {}, msg="duplicate token")
class NodePieceTestCase(RepresentationTestCase):
"""General test case for node piece representations."""
cls = pykeen.nn.node_piece.NodePieceRepresentation
num_relations: ClassVar[int] = 7
num_triples: ClassVar[int] = 31
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
kwargs["triples_factory"] = generation.generate_triples_factory(
num_entities=self.max_id,
num_relations=self.num_relations,
num_triples=self.num_triples,
create_inverse_triples=False,
)
# inferred from triples factory
kwargs.pop("max_id")
return kwargs
def test_estimate_diversity(self):
"""Test estimating diversity."""
diversity = self.instance.estimate_diversity()
assert len(diversity.uniques_per_representation) == len(self.instance.base)
assert 0.0 <= diversity.uniques_total <= 1.0
class EvaluationLoopTestCase(GenericTestCase[pykeen.evaluation.evaluation_loop.EvaluationLoop]):
"""Tests for evaluation loops."""
batch_size: int = 2
factory: CoreTriplesFactory
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
self.factory = Nations().training
kwargs["model"] = FixedModel(triples_factory=self.factory)
return kwargs
@torch.inference_mode()
def test_process_batch(self):
"""Test processing a single batch."""
batch = next(iter(self.instance.get_loader(batch_size=self.batch_size)))
self.instance.process_batch(batch=batch)
class EvaluationOnlyModelTestCase(unittest_templates.GenericTestCase[pykeen.models.EvaluationOnlyModel]):
"""Test case for evaluation only models."""
#: The batch size
batch_size: int = 3
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
kwargs = super()._pre_instantiation_hook(kwargs=kwargs)
dataset = Nations()
self.factory = kwargs["triples_factory"] = dataset.training
return kwargs
def _verify(self, scores: torch.FloatTensor):
"""Verify scores."""
def test_score_t(self):
"""Test score_t."""
hr_batch = self.factory.mapped_triples[torch.randint(self.factory.num_triples, size=(self.batch_size,))][:, :2]
scores = self.instance.score_t(hr_batch=hr_batch)
assert scores.shape == (self.batch_size, self.factory.num_entities)
self._verify(scores)
def test_score_h(self):
"""Test score_h."""
rt_batch = self.factory.mapped_triples[torch.randint(self.factory.num_triples, size=(self.batch_size,))][:, 1:]
scores = self.instance.score_h(rt_batch=rt_batch)
assert scores.shape == (self.batch_size, self.factory.num_entities)
self._verify(scores)
class RankBasedMetricTestCase(unittest_templates.GenericTestCase[RankBasedMetric]):
"""A test for rank-based metrics."""
#: the maximum number of candidates
max_num_candidates: int = 17
#: the number of ranks
num_ranks: int = 33
#: the number of samples to use for monte-carlo estimation
num_samples: int = 1_000
#: the number of candidates for each individual ranking task
num_candidates: numpy.ndarray
#: the ranks for each individual ranking task
ranks: numpy.ndarray
def post_instantiation_hook(self) -> None:
"""Generate a coherent rank & candidate pair."""
self.ranks, self.num_candidates = generate_num_candidates_and_ranks(
num_ranks=self.num_ranks,
max_num_candidates=self.max_num_candidates,
seed=42,
)
def test_docdata(self):
"""Test the docdata contents of the metric."""
self.assertTrue(hasattr(self.instance, "increasing"))
self.assertNotEqual(
"", self.cls.__doc__.splitlines()[0].strip(), msg="First line of docstring should not be blank"
)
self.assertIsNotNone(get_docdata(self.instance), msg="No docdata available")
self.assertIsNotNone(getattr_or_docdata(self.cls, "link"))
self.assertIsNotNone(getattr_or_docdata(self.cls, "name"))
self.assertIsNotNone(getattr_or_docdata(self.cls, "description"))
self.assertIsNotNone(self.instance.key)
def _test_call(self, ranks: numpy.ndarray, num_candidates: Optional[numpy.ndarray]):
"""Verify call."""
x = self.instance(ranks=ranks, num_candidates=num_candidates)
# data type
assert isinstance(x, float)
# value range
self.assertIn(x, self.instance.value_range.approximate(epsilon=1.0e-08))
def test_call(self):
"""Test __call__."""
self._test_call(ranks=self.ranks, num_candidates=self.num_candidates)
def test_call_best(self):
"""Test __call__ with optimal ranks."""
self._test_call(ranks=numpy.ones(shape=(self.num_ranks,)), num_candidates=self.num_candidates)
def test_call_worst(self):
"""Test __call__ with worst ranks."""
self._test_call(ranks=self.num_candidates, num_candidates=self.num_candidates)
def test_call_no_candidates(self):
"""Test __call__ without candidates."""
if self.instance.needs_candidates:
raise SkipTest(f"{self.instance} requires candidates.")
self._test_call(ranks=self.ranks, num_candidates=None)
def test_increasing(self):
"""Test correct increasing annotation."""
x, y = (
self.instance(ranks=ranks, num_candidates=self.num_candidates)
for ranks in [
# original ranks
self.ranks,
# better ranks
numpy.clip(self.ranks - 1, a_min=1, a_max=None),
]
)
if self.instance.increasing:
self.assertLessEqual(x, y)
else:
self.assertLessEqual(y, x)
def _test_expectation(self, weights: Optional[numpy.ndarray]):
"""Test the numeric expectation is close to the closed form one."""
try:
closed = self.instance.expected_value(num_candidates=self.num_candidates, weights=weights)
except NoClosedFormError as error:
raise SkipTest("no implementation of closed-form expectation") from error
generator = numpy.random.default_rng(seed=0)
low, simulated, high = self.instance.numeric_expected_value_with_ci(
num_candidates=self.num_candidates,
num_samples=self.num_samples,
generator=generator,
weights=weights,
)
self.assertLessEqual(low, closed)
self.assertLessEqual(closed, high)
def test_expectation(self):
"""Test the numeric expectation is close to the closed form one."""
self._test_expectation(weights=None)
def test_expectation_weighted(self):
"""Test for weighted expectation."""
self._test_expectation(weights=self._generate_weights())
def _test_variance(self, weights: Optional[numpy.ndarray]):
"""Test the numeric variance is close to the closed form one."""
try:
closed = self.instance.variance(num_candidates=self.num_candidates, weights=weights)
except NoClosedFormError as error:
raise SkipTest("no implementation of closed-form variance") from error
# variances are non-negative
self.assertLessEqual(0, closed)
generator = numpy.random.default_rng(seed=0)
low, simulated, high = self.instance.numeric_variance_with_ci(
num_candidates=self.num_candidates,
num_samples=self.num_samples,
generator=generator,
weights=weights,
)
self.assertLessEqual(low, closed)
self.assertLessEqual(closed, high)
def test_variance(self):
"""Test the numeric variance is close to the closed form one."""
self._test_variance(weights=None)
def test_variance_weighted(self):
"""Test the weighted numeric variance is close to the closed form one."""
self._test_variance(weights=self._generate_weights())
def _generate_weights(self):
"""Generate weights."""
if not self.instance.supports_weights:
raise SkipTest(f"{self.instance} does not support weights")
# generate random weights such that sum = n
generator = numpy.random.default_rng(seed=21)
weights = generator.random(size=self.num_candidates.shape)
weights = self.num_ranks * weights / weights.sum()
return weights
def test_different_to_base_metric(self):
"""Check whether the value is different from the base metric (relevant for adjusted metrics)."""
if not isinstance(self.instance, DerivedRankBasedMetric):
self.skipTest("no base metric")
base_instance = rank_based_metric_resolver.make(self.instance.base_cls)
base_factor = 1 if base_instance.increasing else -1
self.assertNotEqual(
self.instance(ranks=self.ranks, num_candidates=self.num_candidates),
base_factor * base_instance(ranks=self.ranks, num_candidates=self.num_candidates),
)
def test_weights_direction(self):
"""Test monotonicity of weighting."""
if not self.instance.supports_weights:
raise SkipTest(f"{self.instance} does not support weights")
# for sanity checking: give the largest weight to best rank => should improve
idx = self.ranks.argmin()
weights = numpy.ones_like(self.ranks, dtype=float)
weights[idx] = 2.0
weighted = self.instance(ranks=self.ranks, num_candidates=self.num_candidates, weights=weights)
unweighted = self.instance(ranks=self.ranks, num_candidates=self.num_candidates, weights=None)
if self.instance.increasing: # increasing = larger is better => weighted should be better
self.assertLessEqual(unweighted, weighted)
else:
self.assertLessEqual(weighted, unweighted)
def test_weights_coherence(self):
"""Test coherence for weighted metrics & metric in repeated array."""
if not self.instance.supports_weights:
raise SkipTest(f"{self.instance} does not support weights")
# generate two versions
generator = numpy.random.default_rng(seed=21)
repeats = generator.integers(low=1, high=10, size=self.ranks.shape)
# 1. repeat each rank/candidate pair a random number of times
repeated_ranks, repeated_num_candidates = [], []
for rank, num_candidates, repeat in zip(self.ranks, self.num_candidates, repeats):
repeated_ranks.append(numpy.full(shape=(repeat,), fill_value=rank))
repeated_num_candidates.append(numpy.full(shape=(repeat,), fill_value=num_candidates))
repeated_ranks = numpy.concatenate(repeated_ranks)
repeated_num_candidates = numpy.concatenate(repeated_num_candidates)
value_repeat = self.instance(ranks=repeated_ranks, num_candidates=repeated_num_candidates, weights=None)
# 2. do not repeat, but assign a corresponding weight
weights = repeats.astype(float)
value_weighted = self.instance(ranks=self.ranks, num_candidates=self.num_candidates, weights=weights)
self.assertAlmostEqual(value_repeat, value_weighted, delta=2)
class MetricResultTestCase(unittest_templates.GenericTestCase[MetricResults]):
"""Test for metric results."""
def test_flat_dict(self):
"""Test to_flat_dict."""
flat_dict = self.instance.to_flat_dict()
# check flatness
self.assertIsInstance(flat_dict, dict)
for key, value in flat_dict.items():
self.assertIsInstance(key, str)
# TODO: does this suffice, or do we really need float as datatype?
self.assertIsInstance(value, (float, int), msg=key)
self._verify_flat_dict(flat_dict)
def _verify_flat_dict(self, flat_dict: Mapping[str, Any]):
pass
class TrainingInstancesTestCase(unittest_templates.GenericTestCase[Instances]):
"""Test for training instances."""
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
self.factory = Nations().training
return {}
@abstractmethod
def _get_expected_length(self) -> int:
raise NotImplementedError
def test_getitem(self):
"""Test __getitem__."""
self.instance: Instances
assert self.instance[0] is not None
def test_len(self):
"""Test __len__."""
self.assertEqual(len(self.instance), self._get_expected_length())
def test_data_loader(self):
"""Test usage with data loader."""
for batch in torch.utils.data.DataLoader(
dataset=self.instance, batch_size=2, shuffle=True, collate_fn=self.instance.get_collator()
):
assert batch is not None
class BatchSLCWATrainingInstancesTestCase(unittest_templates.GenericTestCase[BaseBatchedSLCWAInstances]):
"""Test for batched sLCWA training instances."""
batch_size: int = 2
num_negatives_per_positive: int = 3
kwargs = dict(
batch_size=batch_size,
negative_sampler_kwargs=dict(
num_negs_per_pos=num_negatives_per_positive,
),
)
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
self.factory = Nations().training
kwargs["mapped_triples"] = self.factory.mapped_triples
return kwargs
def test_data_loader(self):
"""Test data loader."""
for batch in torch.utils.data.DataLoader(dataset=self.instance, batch_size=None):
assert isinstance(batch, SLCWABatch)
assert batch.positives.shape == (self.batch_size, 3)
assert batch.negatives.shape == (self.batch_size, self.num_negatives_per_positive, 3)
assert batch.masks is None
def test_length(self):
"""Test length."""
assert len(self.instance) == len(list(iter(self.instance)))
def test_data_loader_multiprocessing(self):
"""Test data loader with multiple workers."""
self.assertEqual(
sum(
batch.positives.shape[0]
for batch in torch.utils.data.DataLoader(dataset=self.instance, batch_size=None, num_workers=2)
),
self.factory.num_triples,
)
class TrainingCallbackTestCase(unittest_templates.GenericTestCase[TrainingCallback]):
"""Base test case for training callbacks."""
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102
kwargs = super()._pre_instantiation_hook(kwargs)
self.dataset = Nations()
return kwargs
def test_pipeline(self):
"""Test running a small pipeline with the provided callback."""
pipeline(
dataset=self.dataset,
model="distmult",
training_kwargs=dict(
callbacks=self.instance,
),
)
class GraphPairCombinatorTestCase(unittest_templates.GenericTestCase[GraphPairCombinator]):
"""Base test for graph pair combination methods."""
def _add_labels(self, tf: CoreTriplesFactory) -> TriplesFactory:
"""Add artificial labels to a triples factory."""
entity_to_id = {f"e_{i}": i for i in range(tf.num_entities)}
relation_to_id = {f"r_{i}": i for i in range(tf.num_relations)}
return TriplesFactory(
mapped_triples=tf.mapped_triples, entity_to_id=entity_to_id, relation_to_id=relation_to_id
)
def _test_combination(self, labels: bool):
# generate random triples factories
left, right = (generation.generate_triples_factory(random_state=random_state) for random_state in (0, 1))
# generate random alignment
left_idx, right_idx = torch.stack([torch.arange(left.num_entities), torch.randperm(left.num_entities)])[
: left.num_entities // 2
].numpy()
# add label information if necessary
if labels:
left, right = (self._add_labels(tf) for tf in (left, right))
left_idx = [left.entity_id_to_label[i] for i in left_idx]
right_idx = [right.entity_id_to_label[i] for i in right_idx]
# prepare alignment data frame
alignment = pandas.DataFrame(data={EA_SIDE_LEFT: left_idx, EA_SIDE_RIGHT: right_idx})
# call
tf_both, alignment_t = self.instance(left=left, right=right, alignment=alignment)
# check
assert type(tf_both) is type(left)
assert alignment_t.ndim == 2
assert alignment_t.shape[0] == 2
assert alignment_t.shape[1] <= len(alignment)
def test_combination_label(self):
"""Test combination with labels."""
self._test_combination(labels=True)
def test_combination_id(self):
"""Test combination without labels."""
self._test_combination(labels=False)
def test_manual(self):
"""
Smoke-test on a manual example.
cf. https://github.com/pykeen/pykeen/pull/893#discussion_r861553903
"""
left_tf = TriplesFactory.from_labeled_triples(
pandas.DataFrame(
[
["la", "0", "lb"],
["lb", "0", "lc"],
["la", "1", "ld"],
["le", "1", "lg"],
["lh", "1", "lg"],
],
).values
)
right_tf = TriplesFactory.from_labeled_triples(
pandas.DataFrame(
[
["ra", "2", "rb"],
["ra", "2", "rc"],
["rc", "3", "rd"],
["re", "3", "rg"],
["rh", "3", "rg"],
],
).values
)
test_links = pandas.DataFrame(
[
["ld", "rd"],
["le", "re"],
["lg", "rg"],
["lh", "rh"],
],
columns=[EA_SIDE_LEFT, EA_SIDE_RIGHT],
)
combined_tf, alignment_t = self.instance(left=left_tf, right=right_tf, alignment=test_links)
self._verify_manual(combined_tf=combined_tf)
@abstractmethod
def _verify_manual(self, combined_tf: CoreTriplesFactory):
"""Verify the result of the combination of the manual example."""
class EarlyStopperTestCase(unittest_templates.GenericTestCase[EarlyStopper]):
"""Base test for early stopper."""
cls = EarlyStopper
#: The window size used by the early stopper
patience: int = 2
#: The mock losses the mock evaluator will return
mock_losses: list[float] = [10.0, 9.0, 8.0, 9.0, 8.0, 8.0]
#: The (zeroed) index - 1 at which stopping will occur
stop_constant: int = 4
#: The minimum improvement
delta: float = 0.0
#: The best results
best_results: list[float] = [10.0, 9.0, 8.0, 8.0, 8.0]
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
kwargs = super()._pre_instantiation_hook(kwargs)
nations = Nations()
kwargs.update(
dict(
evaluator=MockEvaluator(
key=("hits_at_10", SIDE_BOTH, RANK_REALISTIC),
values=self.mock_losses,
# Set automatic_memory_optimization to false for tests
automatic_memory_optimization=False,
),
model=FixedModel(triples_factory=nations.training),
training_triples_factory=nations.training,
evaluation_triples_factory=nations.validation,
patience=self.patience,
relative_delta=self.delta,
larger_is_better=False,
best_model_path=pathlib.Path(tempfile.gettempdir(), "test-best-model-weights.pt"),
)
)
return kwargs
def test_initialization(self):
"""Test warm-up phase."""
for epoch in range(self.patience):
should_stop = self.instance.should_stop(epoch=epoch)
assert not should_stop
def test_result_processing(self):
"""Test that the mock evaluation of the early stopper always gives the right loss."""
for epoch in range(len(self.mock_losses)):
# Step early stopper
should_stop = self.instance.should_stop(epoch=epoch)
if should_stop:
break
else:
# check storing of results
assert self.instance.results == self.mock_losses[: epoch + 1]
assert self.instance.best_metric == self.best_results[epoch]
def test_should_stop(self):
"""Test that the stopper knows when to stop."""
for epoch in range(self.stop_constant):
self.assertFalse(self.instance.should_stop(epoch=epoch))
self.assertTrue(self.instance.should_stop(epoch=self.stop_constant))
def test_result_logging(self):
"""Test whether result logger is called properly."""
self.instance.result_tracker = mock_tracker = Mock()
self.instance.should_stop(epoch=0)
log_metrics = mock_tracker.log_metrics
self.assertIsInstance(log_metrics, Mock)
log_metrics.assert_called_once()
_, call_args = log_metrics.call_args_list[0]
self.assertIn("step", call_args)
self.assertEqual(0, call_args["step"])
self.assertIn("prefix", call_args)
self.assertEqual("validation", call_args["prefix"])
def test_serialization(self):
"""Test for serialization."""
summary = self.instance.get_summary_dict()
new_stopper = EarlyStopper(
# not needed for test
model=...,
evaluator=...,
training_triples_factory=...,
evaluation_triples_factory=...,
)
new_stopper._write_from_summary_dict(**summary)
for key in summary.keys():
assert getattr(self.instance, key) == getattr(new_stopper, key)
class CombinationTestCase(unittest_templates.GenericTestCase[pykeen.nn.combination.Combination]):
"""Test for combinations."""
input_dims: Sequence[Sequence[int]] = [[5, 7], [5, 7, 11]]
def _iter_input_shapes(self) -> Iterable[Sequence[tuple[int, ...]]]:
"""Iterate over test input shapes."""
for prefix_shape in [tuple(), (2,), (2, 3)]:
for input_dims in self.input_dims:
yield [prefix_shape + (input_dim,) for input_dim in input_dims]
def _create_input(self, input_shapes: Sequence[tuple[int, ...]]) -> Sequence[torch.FloatTensor]:
return [torch.empty(size=size) for size in input_shapes]
def test_inputs(self):
"""Test that the test uses at least one input shape."""
assert list(self._iter_input_shapes())
def test_forward(self):
"""Test forward call."""
for input_shapes in self._iter_input_shapes():
xs = self._create_input(input_shapes=input_shapes)
# verify that the input is valid
assert len(xs) == len(input_shapes)
assert all(x.shape == shape for x, shape in zip(xs, input_shapes))
# combine
x = self.instance(xs=xs)
self.assertIsInstance(x, torch.Tensor)
# verify shape
output_shape = self.instance.output_shape(input_shapes)
self.assertTupleEqual(x.shape, output_shape)
class TextEncoderTestCase(unittest_templates.GenericTestCase[pykeen.nn.text.TextEncoder]):
"""Base tests for text encoders."""
def test_encode(self):
"""Test encoding of texts."""
labels = ["A first sentence", "some other label"]
x = self.instance.encode_all(labels=labels)
assert torch.is_tensor(x)
assert x.shape[0] == len(labels)
class PredictionTestCase(unittest_templates.GenericTestCase[pykeen.predict.Predictions]):
"""Tests for prediction post-processing."""
# to be initialized in subclass
df: pandas.DataFrame
def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
kwargs = super()._pre_instantiation_hook(kwargs)
self.dataset = Nations()
kwargs["factory"] = self.dataset.training
return kwargs
def test_contains(self):
"""Test contains method."""
pred_annotated = self.instance.add_membership_columns(**self.dataset.factory_dict)
assert isinstance(pred_annotated, pykeen.predict.Predictions)
df_annot = pred_annotated.df
# no column has been removed
assert set(df_annot.columns).issuperset(self.df.columns)
# all old columns are unmodified
for col in self.df.columns:
assert (df_annot[col] == self.df[col]).all()
# new columns are boolean
for new_col in set(df_annot.columns).difference(self.df.columns):
assert df_annot[new_col].dtype == bool
def test_filter(self):
"""Test filter method."""
pred_filtered = self.instance.filter_triples(*self.dataset.factory_dict.values())
assert isinstance(pred_filtered, pykeen.predict.Predictions)
df_filtered = pred_filtered.df
# no columns have been added
assert set(df_filtered.columns) == set(self.df.columns)
# check subset relation
assert set(df_filtered.itertuples()).issubset(self.df.itertuples())
class ScoreConsumerTests(unittest_templates.GenericTestCase[pykeen.predict.ScoreConsumer]):
"""Tests for score consumers."""
batch_size: int = 2
num_entities: int = 3
target: Target = LABEL_TAIL
def test_consumption(self):
"""Test calling."""
generator = torch.manual_seed(seed=42)
batch = torch.randint(self.num_entities, size=(self.batch_size, 2), generator=generator)
scores = torch.rand(self.batch_size, self.num_entities)
self.instance(batch=batch, target=self.target, scores=scores)
self.check()
def check(self):
"""Perform additional verification."""
pass