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

from absl.testing import absltest
import numpy as np
from wfa_planning_evaluation_framework.models.reach_point import ReachPoint
from wfa_planning_evaluation_framework.models.reach_surface import ReachSurface
from wfa_planning_evaluation_framework.models.restricted_pairwise_union_reach_surface import (
    RestrictedPairwiseUnionReachSurface,
)
from wfa_planning_evaluation_framework.models.reach_curve import ReachCurve

class LinearCappedReachCurve(ReachCurve):
    """Linear ReachCurve that is capped at max_reach."""

    def _fit(self) -> None:
        self._max_reach = self._data[0].reach()

    def by_impressions(self, impressions: [int], max_frequency: int = 1) -> ReachPoint:
        """Returns the estimated reach for a given impression vector."""
        kplus_frequencies = [
            min(sum(impressions), self.max_reach) // i
            for i in range(1, max_frequency + 1)
        ]
        return ReachPoint(impressions, kplus_frequencies)


class RestrictedPairwiseUnionReachSurfaceTest(absltest.TestCase):
    def assertPointsAlmostEqualToPrediction(
        self, surface, reach_points, tolerance=0.03
    ):
        for reach_point in reach_points:
            self.assertAlmostEqual(
                surface.by_impressions(reach_point.impressions).reach(),
                reach_point.reach(),
                delta=reach_point.reach() * tolerance,
            )

    def generate_true_reach(self, a, reach_curves, impressions):
        p = len(reach_curves)
        reach_vector = [
            reach_curve.by_impressions(impression).reach()
            for reach_curve, impression in zip(reach_curves, impressions)
        ]
        reach = sum(reach_vector)
        for i in range(len(reach_curves)):
            for j in range(len(reach_curves)):
                reach -= (a[i * p + j] * reach_vector[i] * reach_vector[j]) / (
                    max(reach_curves[i].max_reach, reach_curves[j].max_reach) * 2
                )
        return ReachPoint(impressions, [reach])

    def generate_sample_reach_curves(self, num_publishers, decay_rate, universe_size):
        max_reaches = [
            universe_size * (decay_rate ** pub_num) for pub_num in range(num_publishers)
        ]
        return [
            LinearCappedReachCurve([ReachPoint([max_reach], (max_reach,))])
            for max_reach in max_reaches
        ]

    def generate_sample_matrix_a(self, num_publishers):
        lbd = np.array([[2 / num_publishers] for _ in range(num_publishers)])

        true_a = np.ones(num_publishers * num_publishers)
        for i in range(num_publishers):
            for j in range(num_publishers):
                true_a[i * num_publishers + j] = lbd[i] * lbd[j] if i != j else 0
        return true_a

    def generate_sample_reach_points(
        self, true_a, reach_curves, size, universe_size, random_seed
    ):
        reach_points = []
        random_generator = np.random.default_rng(random_seed)
        for _ in range(size):
            impressions = [
                [random_generator.uniform(0, universe_size / 2)]
                for _ in range(len(reach_curves))
            ]
            reach_points.append(
                self.generate_true_reach(true_a, reach_curves, impressions)
            )
        return reach_points

    def test_by_impressions(self):
        num_publishers = 3
        training_size = 50
        universe_size = 200000
        decay_rate = 0.8

        reach_curves = self.generate_sample_reach_curves(
            num_publishers, decay_rate, universe_size
        )
        true_a = self.generate_sample_matrix_a(num_publishers)
        training_reach_points = self.generate_sample_reach_points(
            true_a, reach_curves, training_size, universe_size, 1
        )

        surface = RestrictedPairwiseUnionReachSurface(
            reach_curves, training_reach_points<commit_msg>rebased with fixed pairwise model<commit_after># Copyright 2021 The Private Cardinality Estimation Framework Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for restricted_pairwise_union_reach_surface.py."""

from absl.testing import absltest
import numpy as np
from wfa_planning_evaluation_framework.models.reach_point import ReachPoint
from wfa_planning_evaluation_framework.models.reach_surface import ReachSurface
from wfa_planning_evaluation_framework.models.restricted_pairwise_union_reach_surface import (
    RestrictedPairwiseUnionReachSurface,
)
from wfa_planning_evaluation_framework.models.reach_curve import ReachCurve


class LinearCappedReachCurve(ReachCurve):
    """Linear ReachCurve that is capped at max_reach."""

    def _fit(self) -> None:
        self._max_reach = self._data[0].reach()

    def by_impressions(self, impressions: [int], max_frequency: int = 1) -> ReachPoint:
        """Returns the estimated reach for a given impression vector."""
        kplus_frequencies = [
            min(sum(impressions), self.max_reach) // i
            for i in range(1, max_frequency + 1)
        ]
        return ReachPoint(impressions, kplus_frequencies)


class RestrictedPairwiseUnionReachSurfaceTest(absltest.TestCase):
    def assertPointsAlmostEqualToPrediction(
        self, surface, reach_points, tolerance=0.03
    ):
        for reach_point in reach_points:
            prediction = (
                surface.by_spend(reach_point.spends)
                if (reach_point.spends)
                else surface.by_impressions(reach_point.impressions)
            )
            self.assertAlmostEqual(
                prediction.reach(),
                reach_point.reach(),
                delta=reach_point.reach() * tolerance,
            )

    def generate_true_reach(self, a, reach_curves, impressions, spends=None):
        p = len(reach_curves)
        reach_vector = [
            reach_curve.by_impressions([impression]).reach()
            for reach_curve, impression in zip(reach_curves, impressions)
        ]
        reach = sum(reach_vector)
        for i in range(len(reach_curves)):
            for j in range(len(reach_curves)):
                reach -= (a[i * p + j] * reach_vector[i] * reach_vector[j]) / (
                    max(reach_curves[i].max_reach, reach_curves[j].max_reach) * 2
                )
        return ReachPoint(impressions, [reach], spends)

    def generate_sample_reach_curves(self, num_publishers, decay_rate, universe_size):
        max_reaches = [
            universe_size * (decay_rate ** pub_num) for pub_num in range(num_publishers)
        ]
        return [
            LinearCappedReachCurve([ReachPoint([max_reach], (max_reach,))])
            for max_reach in max_reaches
        ]

    def generate_sample_matrix_a(self, num_publishers):
        lbd = np.array([[2 / num_publishers] for _ in range(num_publishers)])

        true_a = np.ones(num_publishers * num_publishers)
        for i in range(num_publishers):
            for j in range(num_publishers):
                true_a[i * num_publishers + j] = lbd[i] * lbd[j] if i != j else 0
        return true_a

    def generate_sample_reach_points(
        self, true_a, reach_curves, size, universe_size, random_seed
    ):
        reach_points = []
        random_generator = np.random.default_rng(random_seed)
        for _ in range(size):
            impressions = [
                random_generator.uniform(0, universe_size / 2)
                for _ in range(len(reach_curves))
            ]
            reach_points.append(
                self.generate_true_reach(true_a, reach_curves, impressions)
            )
        return reach_points

    def test_by_impressions(self):
        num_publishers = 3
        training_size = 50
        universe_size = 200000
        decay_rate = 0.8

        reach_curves = self.generate_sample_reach_curves(
            num_publishers, decay_rate, universe_size
        )
        true_a = self.generate_sample_matrix_a(num_publishers)
        training_reach_points = self.generate_sample_reach_points(
            true_a, reach_curves, training_size, universe_size, 1
        )

        surface = RestrictedPairwiseUnionReachSurface(
            reach_curves, training_reach_points
{'fraction_non_alphanumeric': 0.0782152792638562, 'fraction_numerical': 0.006847849347314359, 'mean_word_length': 2.7780921584478575, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '15410', 'n_tokens_mistral': 2784, 'n_tokens_neox': 2676, 'n_words': 649}