diff --git a/sdks/python/apache_beam/ml/anomaly/transforms_test.py b/sdks/python/apache_beam/ml/anomaly/transforms_test.py index 423e51abf635..15ec3e994aec 100644 --- a/sdks/python/apache_beam/ml/anomaly/transforms_test.py +++ b/sdks/python/apache_beam/ml/anomaly/transforms_test.py @@ -32,11 +32,11 @@ import apache_beam as beam from apache_beam.ml.anomaly.aggregations import AnyVote +from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.base import AnomalyPrediction from apache_beam.ml.anomaly.base import AnomalyResult from apache_beam.ml.anomaly.base import EnsembleAnomalyDetector from apache_beam.ml.anomaly.detectors.offline import OfflineDetector -from apache_beam.ml.anomaly.detectors.zscore import ZScore from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.specifiable import Specifiable from apache_beam.ml.anomaly.specifiable import _spec_type_to_subspace @@ -116,65 +116,58 @@ def _keyed_result_is_equal_to( return a[0] == b[0] and _unkeyed_result_is_equal_to(a[1], b[1]) +@specifiable +class _ValueDetector(AnomalyDetector): + """Provides order-independent scores for transform orchestration tests. + + A score of -1 represents detector warmup and is returned as NaN. + """ + def learn_one(self, unused_x: beam.Row) -> None: + pass + + def score_one(self, x: beam.Row) -> float: + value = float(next(iter(x))) + return float('NaN') if value == -1 else value + + class TestAnomalyDetection(unittest.TestCase): class TestData: unkeyed_input = [ - beam.Row(x1=1, x2=4), - beam.Row(x1=2, x2=4), - beam.Row(x1=3, x2=5), - beam.Row(x1=10, x2=4), # outlier in key=1, with respect to x1 - beam.Row(x1=2, x2=10), # outlier in key=1, with respect to x2 - beam.Row(x1=3, x2=4), + beam.Row(x1=1, x2=4, score_x1=-1.0, score_x2=-1.0), + beam.Row(x1=2, x2=4, score_x1=-1.0, score_x2=-1.0), + beam.Row(x1=3, x2=5, score_x1=2.0, score_x2=0.0), + beam.Row(x1=10, x2=4, score_x1=4.0, score_x2=1.0), + beam.Row(x1=2, x2=10, score_x1=1.0, score_x2=3.0), + beam.Row(x1=3, x2=4, score_x1=0.0, score_x2=1.0), + ] + keyed_input = list(zip(itertools.repeat(1), unkeyed_input)) + [ + (2, beam.Row(x1=100, x2=5, score_x1=-1.0, score_x2=-1.0)) ] - keyed_input = list(zip(itertools.repeat(1), - unkeyed_input)) + [(2, beam.Row(x1=100, x2=5))] - zscore_x1_expected_predictions = [ - AnomalyPrediction( - model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3), - AnomalyPrediction( - model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3), + score_x1_expected_predictions = [ AnomalyPrediction( - model_id='zscore_x1', - score=2.1213203435596424, - label=0, - threshold=3), + model_id='value_x1', score=float('NaN'), label=-2, threshold=3), AnomalyPrediction( - model_id='zscore_x1', score=8.0, label=1, threshold=3), + model_id='value_x1', score=float('NaN'), label=-2, threshold=3), + AnomalyPrediction(model_id='value_x1', score=2.0, label=0, threshold=3), + AnomalyPrediction(model_id='value_x1', score=4.0, label=1, threshold=3), + AnomalyPrediction(model_id='value_x1', score=1.0, label=0, threshold=3), + AnomalyPrediction(model_id='value_x1', score=0.0, label=0, threshold=3), AnomalyPrediction( - model_id='zscore_x1', - score=0.4898979485566356, - label=0, - threshold=3), - AnomalyPrediction( - model_id='zscore_x1', - score=0.16452254913212455, - label=0, - threshold=3), - AnomalyPrediction( - model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3), + model_id='value_x1', score=float('NaN'), label=-2, threshold=3), ] - zscore_x2_expected_predictions = [ - AnomalyPrediction( - model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2), - AnomalyPrediction( - model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2), - AnomalyPrediction(model_id='zscore_x2', score=0, label=0, threshold=2), + score_x2_expected_predictions = [ AnomalyPrediction( - model_id='zscore_x2', - score=0.5773502691896252, - label=0, - threshold=2), + model_id='value_x2', score=float('NaN'), label=-2, threshold=2), AnomalyPrediction( - model_id='zscore_x2', score=11.5, label=1, threshold=2), + model_id='value_x2', score=float('NaN'), label=-2, threshold=2), + AnomalyPrediction(model_id='value_x2', score=0, label=0, threshold=2), + AnomalyPrediction(model_id='value_x2', score=1.0, label=0, threshold=2), + AnomalyPrediction(model_id='value_x2', score=3.0, label=1, threshold=2), + AnomalyPrediction(model_id='value_x2', score=1.0, label=0, threshold=2), AnomalyPrediction( - model_id='zscore_x2', - score=0.5368754921931594, - label=0, - threshold=2), - AnomalyPrediction( - model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2), + model_id='value_x2', score=float('NaN'), label=-2, threshold=2), ] aggregated_expected_predictions = [ @@ -187,29 +180,29 @@ class TestData: AnomalyPrediction(model_id="custom", label=-2), ] - keyed_zscore_x1_expected = [ + keyed_single_expected = [ (input[0], AnomalyResult(example=input[1], predictions=[decision])) - for input, decision in zip(keyed_input, zscore_x1_expected_predictions) + for input, decision in zip(keyed_input, score_x1_expected_predictions) ] - unkeyed_zscore_x1_expected = [ + unkeyed_single_expected = [ AnomalyResult(example=input, predictions=[decision]) for input, decision in zip( - unkeyed_input, zscore_x1_expected_predictions) + unkeyed_input, score_x1_expected_predictions) ] keyed_ensemble_expected = [( input[0], AnomalyResult(example=input[1], predictions=[decision1, decision2])) for input, decision1, decision2 in zip( - keyed_input, zscore_x1_expected_predictions, - zscore_x2_expected_predictions)] + keyed_input, score_x1_expected_predictions, + score_x2_expected_predictions)] unkeyed_ensemble_expected = [ AnomalyResult(example=input, predictions=[decision1, decision2]) for input, decision1, decision2 in zip( - unkeyed_input, zscore_x1_expected_predictions, - zscore_x2_expected_predictions) + unkeyed_input, score_x1_expected_predictions, + score_x2_expected_predictions) ] keyed_ensemble_agg_expected = [ @@ -225,11 +218,14 @@ class TestData: ] @parameterized.expand([ - (TestData.keyed_input, TestData.keyed_zscore_x1_expected), - (TestData.unkeyed_input, TestData.unkeyed_zscore_x1_expected), + (TestData.keyed_input, TestData.keyed_single_expected), + (TestData.unkeyed_input, TestData.unkeyed_single_expected), ]) def test_one_detector(self, input, expected): - detector = ZScore(features=["x1"], model_id="zscore_x1") + detector = _ValueDetector( + features=["score_x1"], + threshold_criterion=FixedThreshold(3), + model_id="value_x1") with TestPipeline() as p: result = (p | beam.Create(input) | AnomalyDetection(detector)) @@ -244,12 +240,16 @@ def test_one_detector(self, input, expected): ]) def test_multiple_detectors_without_aggregation(self, input, expected): sub_detectors = [] - sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1")) sub_detectors.append( - ZScore( - features=["x2"], + _ValueDetector( + features=["score_x1"], + threshold_criterion=FixedThreshold(3), + model_id="value_x1")) + sub_detectors.append( + _ValueDetector( + features=["score_x2"], threshold_criterion=FixedThreshold(2), - model_id="zscore_x2")) + model_id="value_x2")) with beam.Pipeline() as p: result = ( @@ -267,12 +267,16 @@ def test_multiple_detectors_without_aggregation(self, input, expected): ]) def test_multiple_sub_detectors_with_aggregation(self, input, expected): sub_detectors = [] - sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1")) sub_detectors.append( - ZScore( - features=["x2"], + _ValueDetector( + features=["score_x1"], + threshold_criterion=FixedThreshold(3), + model_id="value_x1")) + sub_detectors.append( + _ValueDetector( + features=["score_x2"], threshold_criterion=FixedThreshold(2), - model_id="zscore_x2")) + model_id="value_x2")) with beam.Pipeline() as p: result = (