wehub-resource-sync
115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
import functools
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Remove the current working directory from sys.path to ensure tests import the
|
|
# installed shap package (with compiled C extensions) rather than the source tree.
|
|
# If this line is commented out, run pytest via `python -P -m pytest tests` instead.
|
|
sys.path[:] = [p for p in sys.path if p not in ("", ".")]
|
|
|
|
try:
|
|
# On MacOS, the newer libomp versions that comes with Homebrew (version >= 12)
|
|
# cause segfaults to occur when pytorch + lightgbm are imported (in that order).
|
|
# The error does not occur when we import lightgbm first because lightgbm
|
|
# distributes its own libomp which takes precedence.
|
|
# cf. GH #3092 for more context.
|
|
import lightgbm # noqa: F401
|
|
except ImportError:
|
|
pass
|
|
|
|
import matplotlib.pyplot as plt # noqa: E402
|
|
import numpy as np # noqa: E402
|
|
import pytest # noqa: E402
|
|
|
|
|
|
def pytest_addoption(parser):
|
|
parser.addoption("--random-seed", action="store", help="Fix the random seed")
|
|
|
|
|
|
@pytest.fixture()
|
|
def random_seed(request) -> int:
|
|
"""Provides a test-specific random seed for reproducible "fuzz testing".
|
|
|
|
Example use in a test:
|
|
|
|
def test_thing(random_seed):
|
|
|
|
# Numpy
|
|
rs = np.random.RandomState(seed=random_seed)
|
|
values = rs.randint(...)
|
|
|
|
# Pytorch
|
|
torch.manual_seed(random_seed)
|
|
|
|
# Tensorflow
|
|
tf.compat.v1.random.set_random_seed(random_seed)
|
|
|
|
By default, a new seed is generated on each run of the tests. If a test
|
|
fails, the random seed used will be displayed in the pytest logs.
|
|
|
|
The seed can be fixed by providing a CLI option e.g:
|
|
|
|
pytest --random-seed 123
|
|
|
|
For numpy usage, note the legacy `RandomState` has stricter version-to-version
|
|
compatibility guarantees than new-style `default_rng`:
|
|
https://numpy.org/doc/stable/reference/random/compatibility.html
|
|
|
|
"""
|
|
manual_seed = request.config.getoption("--random-seed")
|
|
if manual_seed is not None:
|
|
return int(manual_seed)
|
|
else:
|
|
# Otherwise, create a new seed for each test
|
|
rs = np.random.RandomState()
|
|
return rs.randint(0, 1000)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def global_random_seed():
|
|
"""Set the global numpy random seed before each test
|
|
|
|
Nb. Tests that use random numbers should instantiate a local
|
|
`np.random.RandomState` rather than use the global numpy random state.
|
|
"""
|
|
np.random.seed(0)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mpl_test_cleanup():
|
|
"""Run tests in a mpl context manager and close figures after each test."""
|
|
plt.switch_backend("Agg") # Non-interactive backend
|
|
with plt.rc_context():
|
|
yield
|
|
plt.close("all")
|
|
|
|
|
|
def compare_numpy_outputs_against_baseline(*, func_file, baseline_dir=None, rtol=1e-4, atol=1e-6):
|
|
if baseline_dir is None:
|
|
baseline_dir = Path(__file__).parent / "shap_values_baselines"
|
|
elif isinstance(baseline_dir, str):
|
|
baseline_dir = Path(baseline_dir)
|
|
|
|
def decorator(func):
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
output = func(*args, **kwargs)
|
|
base_func_name = f"{Path(func_file).stem}_{func.__name__}"
|
|
baseline_file = baseline_dir / f"{base_func_name}_baseline.npz"
|
|
if hasattr(output, "values"):
|
|
arrays = {"values": output.values, "base_values": np.asarray(output.base_values)}
|
|
else:
|
|
arrays = {"values": output}
|
|
if baseline_file.exists():
|
|
baseline = np.load(baseline_file, allow_pickle=False)
|
|
for key in arrays:
|
|
np.testing.assert_allclose(arrays[key], baseline[key], rtol=rtol, atol=atol)
|
|
else:
|
|
baseline_dir.mkdir(parents=True, exist_ok=True)
|
|
np.savez(baseline_file, **arrays)
|
|
return output
|
|
|
|
return wrapper
|
|
|
|
return decorator
|