src/bin/config_parser/src/privacy/accounting_test.go

// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package privacy

import (
	"math"
	"testing"
)

// equalWithError compares whether two values are at most error apart
func equalWithError(value1, value2, error float64) bool {
	return math.Abs(value1-value2) <= error
}

// arraysEqual compares if two arrays have the same values.
func arraysEqual(array1 []float64, array2 []float64) bool {
	// Check if the arrays have the same length.
	if len(array1) != len(array2) {
		return false
	}
	// Iterate over the arrays and compare each element.
	for i := range array1 {
		if array1[i] != array2[i] {
			return false
		}
	}
	return true
}

// deepEqualWithError compares two maps that define distributions such that for each value
// corresponding probabilities are compared up to specified additive error.
func deepEqualWithError(distr1, distr2 ProbabilityMassFunction, error float64) bool {
	for key, value := range distr1 {
		if !equalWithError(value, distr2[key], error) {
			return false
		}
	}
	for key, value := range distr2 {
		if !equalWithError(value, distr1[key], error) {
			return false
		}
	}
	return true
}

func TestComputeTruncationValueBounds(t *testing.T) {
	distrArray := []float64{0.1, 0.5, 0.41, 0.07, 0.02}
	truncatedMass := 0.2
	minBound, maxBound := computeTruncationValueBounds(distrArray, truncatedMass)
	expectedMinBound, expectedMaxBound := 1, 2
	if minBound != expectedMinBound || maxBound != expectedMaxBound {
		t.Errorf(
			"Incorrect result returned by computeTruncationValueBounds: expected (%d, %d), got (%d, %d)",
			expectedMinBound, expectedMaxBound, minBound, maxBound)
	}
}

func TestComputeTruncationValueBoundsTruncateAll(t *testing.T) {
	distrArray := []float64{0.1, 0.1, 0.1, 0.1, 0.1}
	truncatedMass := 0.600001
	minBound, maxBound := computeTruncationValueBounds(distrArray, truncatedMass)
	expectedMinBound, expectedMaxBound := 3, 1
	if minBound != expectedMinBound || maxBound != expectedMaxBound {
		t.Errorf(
			"Incorrect result returned by computeTruncationValueBounds: expected (%d, %d), got (%d, %d)",
			expectedMinBound, expectedMaxBound, minBound, maxBound)
	}
}

func TestMinMaxKeyValues(t *testing.T) {
	pmf := ProbabilityMassFunction{
		-1: 0.1,
		5:  0.2,
		2:  0.4,
	}
	minKey, maxKey := minMaxKeyValues(pmf)
	expectedMinKey, expectedMaxKey := -1, 5
	if minKey != expectedMinKey || maxKey != expectedMaxKey {
		t.Errorf(
			"Incorrect result returned by minMaxKeyValue: expected (%d, %d), got (%d, %d)", expectedMinKey, expectedMaxKey, minKey, maxKey)
	}
}

func TestConvertMapToShiftedArrayError(t *testing.T) {
	pmf := ProbabilityMassFunction{
		-1: 0.3,
		1:  0.3,
		3:  0.4,
	}
	size := 3
	_, err := convertMapToShiftedArray(pmf, size)

	if err == nil {
		t.Errorf("Expected error, but no error was found.")
	}
}

func TestConvertMapToShiftedArray(t *testing.T) {
	pmf := ProbabilityMassFunction{
		-1: 0.3,
		1:  0.3,
		3:  0.4,
	}
	size := 6
	shiftedDistr, _ := convertMapToShiftedArray(pmf, size)

	array := []float64{0.3, 0.0, 0.3, 0.0, 0.4, 0.0}
	expectedShiftedDistr := shiftedDistribution{array, -1}
	if !arraysEqual(shiftedDistr.array, expectedShiftedDistr.array) ||
		shiftedDistr.shift != expectedShiftedDistr.shift {
		t.Errorf(
			"Incorrect result of convertMapToShiftedArray. Expected %v, got %v",
			expectedShiftedDistr, shiftedDistr)
	}
}

func TestConvolveHeterogeneousDistributionsNoTruncation(t *testing.T) {
	pmf1 := ProbabilityMassFunction{0: 0.5, 1: 0.5, 2: 0.0}
	pmf2 := ProbabilityMassFunction{0: 0.0, 1: 0.5, 2: 0.5}
	result, err := convolveHeterogeneousDistributions(pmf1, pmf2, 0)
	if err != nil {
		t.Errorf("ConvolveHeterogeneousDistributions returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{
		0: 0.0,
		1: 0.25,
		2: 0.5,
		3: 0.25,
		4: 0.0,
	}

	if !deepEqualWithError(expectedPmf, result, 0.00001) {
		t.Errorf("Incorrect result of convolution of a pair of distributions: expected %v, got %v", expectedPmf, result)
	}
}

func TestConvolveHeterogeneousDistributionsWithTruncation(t *testing.T) {
	pmf1 := ProbabilityMassFunction{0: 0.1, 1: 0.9, 2: 0.0}
	pmf2 := ProbabilityMassFunction{0: 0.0, 1: 0.9, 2: 0.1}
	result, err := convolveHeterogeneousDistributions(pmf1, pmf2, 0.2)
	if err != nil {
		t.Errorf("ConvolveHeterogeneousDistributions returned error: %s", err.Error())
	}
	// Mass of 0.09 at 1 and 3 respectively get truncated.
	expectedPmf := ProbabilityMassFunction{2: 0.82}
	if !deepEqualWithError(expectedPmf, result, 0.00001) {
		t.Errorf("Incorrect result of convolution of a pair of distributions: expected %v, got %v", expectedPmf, result)
	}
}

func TestConvolve1Time(t *testing.T) {
	pmf := ProbabilityMassFunction{0: 0.5, 1: 0.5}
	result, err := convolveHomogeneousDistributions(pmf, 1, 0)
	if err != nil {
		t.Errorf("ConvolveHomogeneousDistributions returned error: %s", err.Error())
	}
	if !deepEqualWithError(pmf, result, 0.00001) {
		t.Errorf("Incorrect result of 1 convolution computation: expected %v, got %v", pmf, result)
	}
}

func TestConvolve3Times(t *testing.T) {
	pmf := ProbabilityMassFunction{0: 0.5, 1: 0.5}
	result, err := convolveHomogeneousDistributions(pmf, 3, 1e-5)
	if err != nil {
		t.Errorf("ConvolveHomogeneousDistributions returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{
		0: 0.125,
		1: 0.375,
		2: 0.375,
		3: 0.125,
	}
	if !deepEqualWithError(expectedPmf, result, 1e-10) {
		t.Errorf("Incorrect result of 3-convolution computation: expected %v, got %v", expectedPmf, result)
	}
}

func TestConvolve1000Times(t *testing.T) {
	pmf := ProbabilityMassFunction{1: 1.0}
	result, err := convolveHomogeneousDistributions(pmf, 1000, 1e-5)
	if err != nil {
		t.Errorf("ConvolveHomogeneousDistributions returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{1000: 1.0}

	if !deepEqualWithError(expectedPmf, result, 1e-10) {
		t.Errorf("Incorrect result of 1000-convolution computation: expected %v, got %v", expectedPmf, result)
	}
}

func TestComputeKConvolutionTruncationParams(t *testing.T) {
	distrArray := []float64{0.1, 0.4, 0.5}
	minBound, maxBound := computeTruncationValueBoundsForHomogeneousConvolution(distrArray, 3, 0.5)
	if minBound != 2 || maxBound != 6 {
		t.Errorf("Incorrect result of computation truncation parameters for KConvolution: expected 2 and 6, got %d and %d", minBound, maxBound)
	}
}

func TestRescaleFFTArray(t *testing.T) {
	distr := []float64{3, 2}
	result := rescaleFFTArray(distr)
	if result[0] != 1.5 || result[1] != 1 {
		t.Errorf("Incorrect result of rescaling of a distribution: expected [0:1.5 1:1], got %v", result)
	}
}

func TestGeneratePLDPessimisticZeroInfiniteMass(t *testing.T) {
	pmfUpper := ProbabilityMassFunction{1: 0.2, 2: 0.8}
	pmfLower := ProbabilityMassFunction{1: 0.5, 2: 0.5}
	result := generatePLD(pmfUpper, pmfLower, -50, 1e-4, true)
	expectedPmf := ProbabilityMassFunction{-9162: 0.2, 4701: 0.8}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.00001) {
		t.Errorf("Incorrect result of PLD generation: expected %v got %v", expectedPmf, result.pmf)
	}
	if result.infiniteMass != 0.0 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 0.0 got %f", result.infiniteMass)
	}
}

func TestGeneratePLDOptimistic(t *testing.T) {
	pmfUpper := ProbabilityMassFunction{1: 0.6, 2: 0.4}
	pmfLower := ProbabilityMassFunction{1: 0.5, 2: 0.5}
	result := generatePLD(pmfUpper, pmfLower, -0.55, 1e-4, false)
	expectedPmf := ProbabilityMassFunction{1823: 0.6}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.00001) {
		t.Errorf("Incorrect result of PLD generation: expected %v got %v", expectedPmf, result.pmf)
	}
	if result.infiniteMass != 0.0 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 0.0 got %f", result.infiniteMass)
	}
}

func TestGeneratePLDSmallandLargeDiscretization(t *testing.T) {
	pmfUpper := ProbabilityMassFunction{1: 0.7, 2: 0.3}
	pmfLower := ProbabilityMassFunction{1: 0.5, 2: 0.5}

	// Small discretization : 1e-6
	result := generatePLD(pmfUpper, pmfLower, -0.55, 1e-6, false)
	expectedPmf := ProbabilityMassFunction{336472: 0.7}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.00001) {
		t.Errorf("Incorrect result of PLD generation: expected %v got %v", expectedPmf, result.pmf)
	}
	if result.infiniteMass != 0.0 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 0.0 got %f", result.infiniteMass)
	}

	// Large discretization : 1e-2
	result = generatePLD(pmfUpper, pmfLower, -0.55, 1e-2, false)
	expectedPmf = ProbabilityMassFunction{33: 0.7}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.00001) {
		t.Errorf("Incorrect result of PLD generation: expected %v got %v", expectedPmf, result.pmf)
	}
	if result.infiniteMass != 0.0 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 0.0 got %f", result.infiniteMass)
	}
}

func TestGenerateEmptyPLDPessimistic(t *testing.T) {
	pmfUpper := ProbabilityMassFunction{1: 1.0}
	pmfLower := ProbabilityMassFunction{2: 0.5, 4: 0.5}
	result := generatePLD(pmfUpper, pmfLower, -50.0, 1e-4, true)
	if len(result.pmf) != 0 {
		t.Errorf("Incorrect result of PLD generation: expected empty distribution got %v", result.pmf)
	}
	if result.infiniteMass != 1.0 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 1.0 got %f", result.infiniteMass)
	}
}

func TestGenerateEmptyPLDOptimistic(t *testing.T) {
	pmfUpper := ProbabilityMassFunction{1: 1.0}
	pmfLower := ProbabilityMassFunction{2: 0.5, 4: 0.5}
	result := generatePLD(pmfUpper, pmfLower, -50.0, 1e-4, false)
	if len(result.pmf) != 0 {
		t.Errorf("Incorrect result of PLD generation: expected empty distribution got %v", result.pmf)
	}
	if result.infiniteMass != 1.0 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 1.0 got %f", result.infiniteMass)
	}
}

func TestGeneratePLDNonzeroInfiniteMass(t *testing.T) {
	pmfUpper := ProbabilityMassFunction{1: 0.6, 2: 0.4}
	pmfLower := ProbabilityMassFunction{1: 0.5, 3: 0.5}
	result := generatePLD(pmfUpper, pmfLower, -50, 1e-4, true)
	expectedPmf := ProbabilityMassFunction{1824: 0.6}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.00001) {
		t.Errorf("Incorrect result of PLD generation: expected %v got %v", expectedPmf, result.pmf)
	}
	if result.infiniteMass != 0.4 {
		t.Errorf("Incorrect result of PLD generation: expected infinite mass 0.0 got %f", result.infiniteMass)
	}
}

func TestComposeHeterogeneousPLD(t *testing.T) {
	discretization := 0.01
	pld1 := PrivacyLossDistribution{
		discretization:        discretization,
		infiniteMass:          0.1,
		pmf:                   ProbabilityMassFunction{0: 0.05, 1: 0.7, 2: 0.15},
		pessimisticEstimation: true}
	pld2 := PrivacyLossDistribution{
		discretization:        discretization,
		infiniteMass:          0.05,
		pmf:                   ProbabilityMassFunction{0: 0.2, 1: 0.65, 2: 0.1},
		pessimisticEstimation: true}
	result, err := composeHeterogeneousPLD(&pld1, &pld2, 0.005)
	if err != nil {
		t.Errorf("ComposeHeterogeneousPLD returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{
		0: 0.01,
		1: 0.1725,
		2: 0.49,
		3: 0.1675,
		4: 0.015,
	}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.000001) {
		t.Errorf("Incorrect result of PLD product: expected %v got %v", expectedPmf, result.pmf)
	}
	if !equalWithError(result.infiniteMass, 0.15, 0.000001) {
		t.Errorf("Incorrect result of PLD product: expected infinite mass 0.15 got %f", result.infiniteMass)
	}
}

func TestComposeHeterogeneousPLDWithTruncation(t *testing.T) {
	discretization := 0.01
	pld1 := PrivacyLossDistribution{
		discretization:        discretization,
		infiniteMass:          0.0,
		pmf:                   ProbabilityMassFunction{0: 0.05, 1: 0.9, 2: 0.05},
		pessimisticEstimation: true}
	pld2 := PrivacyLossDistribution{
		discretization:        discretization,
		infiniteMass:          0.0,
		pmf:                   ProbabilityMassFunction{0: 0.1, 1: 0.8, 2: 0.1},
		pessimisticEstimation: true}
	result, err := composeHeterogeneousPLD(&pld1, &pld2, 0.011)
	if err != nil {
		t.Errorf("ComposeHeterogeneousPLD returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{
		1: 0.13,
		2: 0.73,
		3: 0.13,
	}
	if !deepEqualWithError(expectedPmf, result.pmf, 0.000001) {
		t.Errorf("Incorrect result of PLD product: expected %v got %v", expectedPmf, result.pmf)
	}
	if !equalWithError(result.infiniteMass, 0.011, 0.000001) {
		t.Errorf("Incorrect result of PLD product: expected infinite mass 0.011 got %f", result.infiniteMass)
	}
}

func TestComposeHeterogeneousPLDDifferentDiscretizationError(t *testing.T) {
	pld1 := PrivacyLossDistribution{
		discretization:        0.01,
		infiniteMass:          0.0,
		pmf:                   ProbabilityMassFunction{0: 0.3, 1: 0.7},
		pessimisticEstimation: true}
	pld2 := PrivacyLossDistribution{
		discretization:        0.0001,
		infiniteMass:          0.0,
		pmf:                   ProbabilityMassFunction{0: 0.3, 1: 0.7},
		pessimisticEstimation: true}
	_, err := composeHeterogeneousPLD(&pld1, &pld2, 0.005)
	if err == nil {
		t.Errorf("ComposeHeterogeneousPLD didn't return an error while composing PLDs with distinct discretization")
	}
}

func TestSelfCompositionWithTailTruncation(t *testing.T) {
	pld := PrivacyLossDistribution{
		discretization:        0.1,
		infiniteMass:          0.1,
		pmf:                   ProbabilityMassFunction{0: 0.05, 1: 0.7, 2: 0.15},
		pessimisticEstimation: true}
	truncatedMass := 0.001
	result, err := composeHomogeneousPLD(&pld, 4, truncatedMass)
	if err != nil {
		t.Errorf("ComposeHomogeneousPLD returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{
		1: 0.00035,
		2: 0.007425,
		3: 0.07175,
		4: 0.284538,
		5: 0.21525,
		6: 0.066825,
		7: 0.00945,
		8: 0.000506,
	}
	if !deepEqualWithError(expectedPmf, result.pmf, 1e-5) {
		t.Errorf("Incorrect result of PLD self-composition with tail-mass truncation: expected %v, got %v", expectedPmf, result.pmf)
	}
	if !equalWithError(0.3439+truncatedMass, result.infiniteMass, 0.00001) {
		t.Errorf("Incorrect result of PLD self-composition with tail-mass truncation: expected infinite mass %f got %f", 0.1+truncatedMass, result.infiniteMass)
	}
}

func TestPLDSelfComposition(t *testing.T) {
	pld := PrivacyLossDistribution{
		discretization:        0.1,
		infiniteMass:          0.1,
		pmf:                   ProbabilityMassFunction{0: 0.05, 1: 0.7, 2: 0.15},
		pessimisticEstimation: true}
	result, err := composeHomogeneousPLD(&pld, 4, 0)
	if err != nil {
		t.Errorf("ComposeHomogeneousPLD returned error: %s", err.Error())
	}
	expectedPmf := ProbabilityMassFunction{
		0: 0.000006,
		1: 0.00035,
		2: 0.007425,
		3: 0.07175,
		4: 0.284538,
		5: 0.21525,
		6: 0.066825,
		7: 0.00945,
		8: 0.000506,
	}
	if !deepEqualWithError(expectedPmf, result.pmf, 1e-5) {
		t.Errorf("Incorrect result of PLD self-composition: expected %v, got %v", expectedPmf, result.pmf)
	}
	if !equalWithError(0.3439, result.infiniteMass, 0.0001) {
		t.Errorf("Incorrect result of PLD self-composition: expected infinite mass %f got %f", 0.3439, result.infiniteMass)
	}
}

func TestGetDeltaForPrivacyLossDistribution(t *testing.T) {
	pmf := ProbabilityMassFunction{
		1824:  0.6,
		-2231: 0.4,
	}
	infMass := 0.0

	pld := PrivacyLossDistribution{
		discretization:        0.0001,
		infiniteMass:          infMass,
		pmf:                   pmf,
		pessimisticEstimation: false}

	divergence := pld.getDeltaForPrivacyLossDistribution(0.0)
	if !equalWithError(0.1, divergence, 0.0001) {
		t.Errorf(
			"Incorrect result of getDeltaForPrivacyLossDistribution: expected divergence value to be %f got %f",
			0.1, divergence)
	}
}