noise_corrector.py

#!/usr/bin/env python3

import argparse
import collections
from datetime import timedelta
import json
import numpy
import random
from typing import Any, List, Hashable, Tuple, Dict, TypedDict, TypeVar, Generator, Iterable
import sys

# This file provides functionality to correct noisy data coming from the event-level reports in the Attribution Reporting API.
#
# Noise rate values are taken from the EVENT.md explainer:
# https://github.com/WICG/conversion-measurement-api/blob/main/EVENT.md#data-limits-and-noise

VERSION = "noise_corrector.py 0.0.1"

################## GENERIC UTILITIES ##################

# Useful type declarations for working with JSON:
class SourceRegistrationConfig(TypedDict, total=False):
  source_event_id: str
  expiry: int # Optional


class Source(TypedDict):
  source_time: int
  source_type: str
  registration_config: SourceRegistrationConfig


class ReportRaw(TypedDict):
  source_event_id: str
  source_type: str
  trigger_data: str


class Report(TypedDict):
  report_time: int
  report: ReportRaw


Joined = Tuple[Source, List[Report]]
JoinedList = List[Joined]
JoinedGen  = Generator[Joined, None, None]
# Raw output configurations for correcting randomized response.
# The first int is window_index (i.e. which reporting window the output fell in).
# The second int is trigger_data.
OutputConfig = Tuple[Tuple[int, int], ...]

# For generic functions.
T = TypeVar('T')

OutputEnumerationGen = Generator['OutputEnumeration', None, None]

class ParamConfig:
  def __init__(self,
               randomized_reponse_rate,
               trigger_data_cardinality,
               max_reports,
               max_windows,
               name):
    self.rr_rate = randomized_reponse_rate
    self.data_cardinality = trigger_data_cardinality
    self.max_reports = max_reports
    self.max_windows = max_windows
    self.name = name


# https://github.com/WICG/conversion-measurement-api/blob/main/EVENT.md#data-limits-and-noise
NAV_PARAMS = ParamConfig(.0024, 8, 3, 3, 'navigation')
EVENT_PARAMS = ParamConfig(.0000025, 2, 1, 1, 'event')

# See 9.8.11
# https://wicg.github.io/conversion-measurement-api/#obtaining-attribution-source-params
DEFAULT_EXPIRY = 60 * 60 * 24 * 30

class OutputEnumeration:
  # Class encapsulating logic which represents the logical output of the API
  # per source event, for the purposes of debiasing the randomized response
  # privacy mechanism.
  def __init__(self, report_tuple: OutputConfig, params: ParamConfig):
    '''report_tuple must be in the form ((window_index, trigger_data),...)'''
    self.output: OutputConfig = tuple(sorted(report_tuple))
    self.params = params    

  def __str__(self):
    def report_to_string(r):
      if r is None:
        return 'No report'
      return {'window': r[0], 'data': r[1]}
    return json.dumps([report_to_string(r) for r in self.output])

  def __repr__(self): return self.output.__repr__()
  def __hash__(self): return self.output.__hash__()

  # Needed to support instances in dictionaries
  def __eq__(self, other): return self.output.__eq__(other.output)
  def __lt__(self, other): return self.output.__lt__(other.output)

  @classmethod
  def create_null(cls, params: ParamConfig):
    return cls((), params)

  @classmethod
  def create_from_data(cls, source_with_reports):
    '''Create an output directly from structured data

    source_with_reports: list of (source, [report, ...]) tuples.
    See get_conversion_counts for expected structure for sources and reports.
    '''
    source, reports = source_with_reports
    params = EVENT_PARAMS if source['source_type'] == 'event' else NAV_PARAMS
    max_reports = params.max_reports
    assert len(reports) <= max_reports

    def parse_report(source, report: dict) -> Tuple[int, int]:
      window_index = OutputEnumeration._get_window_index_for_report(
          source, report)
      trigger_data = int(report['report']['trigger_data'])
      return (window_index, trigger_data)

    return cls(tuple(parse_report(source, r) for r in reports), params)


  @classmethod
  def generate_all(cls, params: ParamConfig) -> OutputEnumerationGen:
    """Generator to generate every single possible OutputEnumeration for the given `params`"""
    def gen_all(max_reports, index, output_so_far) -> OutputEnumerationGen:
      """Helper generator which aggregates outputs in `output_so_far` recursively"""
      yield cls(tuple(output_so_far), params)
      if max_reports == 0:
        return

      # Build up `output_so_far` recursively.
      # 1. Insert the report specified by the current index
      # 2. Recurse down to build up the remaining reports conditioned on
      #    indices >= the current index.
      # The condition in step 2 is important because it means we keep a constraint
      # where `output_so_far` is increasing in indices and that we don't generate
      # duplicate outputs like reports indexed by (1,2,3) and (3,2,1).
      for next_index in range(index, params.data_cardinality * params.max_windows):
        window, data = divmod(next_index, params.data_cardinality)
        output_so_far.append((window, data))
        yield from gen_all(max_reports - 1, next_index, output_so_far)
        output_so_far.pop()
    return gen_all(params.max_reports, 0, [])


  def data_histogram(self) -> numpy.ndarray:
    '''Returns the number of reports associated with this output, broken
    out by the trigger_data'''

    return numpy.bincount([o[1] for o in self.output], minlength=self.params.data_cardinality)

  def get_report_time(self, window: int, source: Source) -> int:
    """For a given Source and window index, returns the estimated report time.
    This is used for generating synthetic event-level reports"""
    expiry = source['Attribution-Reporting-Register-Source'].get('expiry', DEFAULT_EXPIRY)
    if self.params.name == 'event':
      return expiry

    # TODO(csharrison): This assumes the expiry > 7 days.
    windows = [timedelta(days=2), timedelta(days=7), timedelta(seconds=expiry)]
    source_time = int(source['timestamp'])
    return int(source_time + windows[window].total_seconds())

  def generate_reports_for_source(self, source: Source) -> List[Report]:
    assert self.params.name == source['source_type']
    source_event_id = source['Attribution-Reporting-Register-Source']['source_event_id']

    reports: List[Report] = []
    for window, trigger_data in self.output:
      reports.append({
          'report_time': str(self.get_report_time(window, source)),
          'report': {
              'source_event_id': source_event_id,
              'source_type': self.params.name,
              'trigger_data': str(trigger_data)
          }
      })
    return reports

  @staticmethod
  def _get_window_index_for_report(source: dict, report: dict) -> int:
    if source['source_type'] == 'event':
      return 0

    report_time = int(report['report_time'])
    source_time = int(source['timestamp'])
    raw_expiry = source['Attribution-Reporting-Register-Source'].get('expiry', DEFAULT_EXPIRY)
    expiry = timedelta(seconds=int(raw_expiry))
    delta = timedelta(seconds=report_time - source_time)

    # Navigation sources have 3 possible reporting windows.
    # TODO(csharrison): this assumes that the expiry > 7 days. This should be
    # updated to be more robust.
    assert expiry > timedelta(days=7)
    if delta >= expiry:
      return 2
    if delta >= timedelta(days=7):
      return 1
    return 0


# Estimates the true values for data generated through k-ary randomized
# response. The estimator should be unbiased.
def estimate_true_values(observed_values: Dict[T, int],
                         randomized_response_rate: float) -> Dict[T, float]:
  '''Returns a list of corrected counts

  observed_values: A map of size `k` for k-ary randomized response. Each item
                   represents a randomized response output as the key, with its
                   associated count as the value.

  randomized_response_rate: The probability any given output is randomized
  '''
  n = sum(observed_values.values())
  k = len(observed_values)
  x = randomized_response_rate

  # Proof: see formula 6 in https://arxiv.org/pdf/1602.07387.pdf. Note in that
  # formulation, `x` = (k / (k + exp(epsilon) - 1)).
  # Performing that substitution yields:
  # estimate(v) = (v(k + exp(eps) - 1) - n) / (exp(eps) - 1)
  # Which matches formula 6 after applying a multiplication of `n` to transform
  # rates to counts.
  def estimate(v):
    return (v - n * x / k) / (1 - x)

  return {k: estimate(v) for k, v in observed_values.items()}


def join_reports_with_sources(sources: List[Source], reports: List[Report]) -> JoinedGen:
  '''Returns tuple of (source, [reports]) joined with the source_event_id'''
  report_map = collections.defaultdict(list)
  for r in reports:
    report_map[r['report']['source_event_id']].append(r)

  def join(source) -> Joined:
    source_event_id = source['Attribution-Reporting-Register-Source']['source_event_id']
    return (source, report_map[source_event_id])
  return (join(s) for s in sources)


################## AGGREGATE UTILITIES ##################

def get_raw_corrected_map(joined: Iterable[Joined], params: ParamConfig) -> Dict[OutputEnumeration, float]:
  '''Returns an aggregate map of OutputEnumeration -> float, with noise debiased
  according to params. Each returned float is an estimate count for how many times the
  given OutputEnumeration showed up in `joined`'''
  output_counts = {o: 0 for o in OutputEnumeration.generate_all(params)}
  for j in joined:
    output = OutputEnumeration.create_from_data(j)
    assert output in output_counts
    output_counts[output] += 1

  return estimate_true_values(output_counts, params.rr_rate)


def correct_aggregates(joined: JoinedGen, params: ParamConfig) -> List[dict]:
  '''Returns debiased report counts broken out by trigger_data'''
  corrected = get_raw_corrected_map(joined, params)
  histogram = numpy.zeros(params.data_cardinality)
  for o, c in corrected.items():
    histogram += c * o.data_histogram()
  x = 10
  return [{'trigger_data': t,
           'report_count': v} for t, v in enumerate(histogram)]


################## EVENT-LEVEL UTILITIES ##################

def adjust_to_match_distribution(values: List[T],
                                 estimated_counts: Dict[T, float],
                                 default_value: T) -> List[T]:
  '''Given a list of values, and an aggregate map that counts all possible values,
  returns a list of adjusted values to match the aggregate distribution.

  values: list of arbitrary objects (must be possible to place in dict/set). To
          minimize bias, this list should be shuffled prior to calling this
          method.
  estimated_counts: map of value -> estimated count. All values in `values` must be
                    present. Note that negatives are allowed, but they will be
                    clamped to 0 internally. Note that because the we cannot
                    assign any output less than 0, the end values will always be
                    biased.
  default_value: a default / null value

  The algorithm used in this method attempts to achieve two goals:
  1. The output values match the distribution as much as possible, minimizing
     bias.
  2. The output maximizes 'overlap' with the input. That is, satisfying (1), we
     want to minimize the number of inputs that are adjusted.
  '''

  # Split the distribution into two chunks: elements with mass >= 1 and not.
  # Clamp negatives to 0. Note this introduces bias, but it seems unavoidable.
  large_elts = {k: v for k, v in estimated_counts.items() if v >= 1}
  small_elts = {k: max(0, v) for k, v in estimated_counts.items() if v < 1}
  assert len(small_elts) + len(large_elts) == len(estimated_counts)

  def decrement_large_elt(elt):
    large_elts[elt] -= 1
    if large_elts[elt] < 1:
      small_elts[elt] = large_elts[elt]
      large_elts.pop(elt)

  def handle_value(val):
    assert val in estimated_counts and (val in small_elts or val in large_elts)
    # Do not adjust s if there is mass in the aggregate distribution for it.
    if val in large_elts:
      decrement_large_elt(val)
      return val

    # Otherwise, pick from among the elements with enough mass for subtraction.
    if len(large_elts) > 0:
      elt = random.choices(list(large_elts.keys()), large_elts.values())[0]
      decrement_large_elt(elt)
      return elt

    # Otherwise, sample with replacement from the small elts, with any leftover
    # mass defaulting to a default value / null value. This also introduces
    # small bias towards the null value bucket.
    total_mass_remaining = sum(small_elts.values())
    if random.random() <= total_mass_remaining:
      return random.choices(list(small_elts.keys()), small_elts.values())[0]
    return default_value

  return [handle_value(v) for v in values]


def generate_corrected_event_level(joined: JoinedList, params: ParamConfig) -> JoinedList:
  # Shuffle the input data. This is because the order of which values are
  # considered in the loop below matters, because we adjust the distribution
  # as we go. Shuffling first avoids favoring some values over others.
  # Note: ideally adjust_to_match_distribution would shuffle, but it becomes
  # more difficult to associate the adjustments with the input JSON.
  random.shuffle(joined)
  aggregates = get_raw_corrected_map(joined, params)
  outputs = [OutputEnumeration.create_from_data(j) for j in joined]
  adjusted_values = adjust_to_match_distribution(
      outputs, aggregates, OutputEnumeration.create_null(params))

  def adjust_json(j, adjusted) -> Joined:
    source: Source = j[0]
    adjusted_reports: List[Report] = adjusted.generate_reports_for_source(
        source)
    return (source, adjusted_reports)

  return [adjust_json(j, a) for j, a in zip(joined, adjusted_values)]


def main():
  DESCRIPTION = '''\
  noise_corrector.py is a command-line tool that attempts to debias the noise
  output in the Attribution Reporting API for event-level reports. It takes as
  input JSON files which represent source registrations and the resulting
  event-level reports. It can provide debiased output in aggregate, or via
  generating synthetic event-level data.

  The noise added in the API is described at:
  https://github.com/WICG/conversion-measurement-api/blob/main/EVENT.md#data-limits-and-noise
  This utility expects JSON input coming from stdin in the following format:
  {
    "input": {
      // List of zero or more sources.
      "sources": [
        {
          // Required time at which to register the source in seconds since the
          // UNIX epoch.
          "timestamp": "123",

          // Required source type, either "navigation" or "event", corresponding to
          // whether the source is registered on click or on view, respectively.
          "source_type": "navigation",

          "Attribution-Reporting-Register-Source": {
            // Required uint64 formatted as a base-10 string.
            "source_event_id": "123456789",

            // Optional int64 in milliseconds formatted as a base-10 string.
            // Defaults to 30 days.
            "expiry": "864000000"
          }
        },
        ...
      ]
    },
    // List of zero or more reports.
    "reports": [
      {
        // Time at which the report would have been sent in seconds since the
        // UNIX epoch.
        "report_time": "123",

        // The report itself.
        "report": {
          "source_event_id": "1337",
          "source_type": "navigation",

          // Coarse data set in the attribution trigger registration
          "trigger_data": "4"
        }
      },
      ...
    ]
  }
  '''
  parser = argparse.ArgumentParser(
      description=DESCRIPTION, formatter_class=argparse.RawTextHelpFormatter)

  INPUT_MODE_HELP = '''\
  Optional. One of `single` or `multi`. Defaults to `single`, which reads all of
  stdin as one JSON file, with the format above. `multi` mode accepts input in
  the JSON-lines format, where every separate line of input follows the above
  format. It is expected in `multi` mode that events from a given browser are
  isolated in a particular line of input (i.e. the tool will treat data on
  different lines completely independently).
  '''
  parser.add_argument('--input_mode', dest='input_mode',
                      choices=['single', 'multi'],
                      default='single', help=INPUT_MODE_HELP)

  OUTPUT_MODE_HELP = '''\
  Optional. One of `experimental-event-level` or `aggregate`. Defaults to
  `aggregate`, which provides aggregate counts for different trigger data
  values. The `experimental-event-level` mode outputs synthetic sources and
  their associated reports matching the debiased aggregate data as best as
  possible. In other words, the tool will mutate the input reports in some way
  (sometimes by generating new reports that were not input into the system) so
  that the output event-level data, when aggregated, matches the data that
  `aggregate` mode would return.

  Note: the `experimental-event-level` format is experimental. More work is
  needed to make sure the use-cases for it are well-understood.

  `aggregate` output format:
  {
    "navigation": [
      {
        "trigger_data": 0,
        "report_count": <count>
      },
    ...
    ],
    "event": [
      {
        "trigger_data": 0,
        "report_count": <count>
      },
      ...
    ]
  }

  `experimental-event-level` output format:
  [
    {
      "source": <source, same as input format>
      "reports": <same as input reports format>
    },
    ...
  ]
  '''
  parser.add_argument('--output_mode', dest='output_mode',
                      choices=['experimental-event-level', 'aggregate'],
                      default='aggregate', help=OUTPUT_MODE_HELP)

  parser.add_argument('-v', '--version', action='version', version=VERSION)

  args = parser.parse_args()

  def gen_joined() -> JoinedGen:
    if args.input_mode == 'single':
      input = json.load(sys.stdin)
      yield from join_reports_with_sources(input['input']['sources'], input['event_level_reports'])
    else:
      for line in sys.stdin:
        input = json.loads(line)
        yield from join_reports_with_sources(input['input']['sources'], input['event_level_reports'])

  joined = gen_joined()
  navs: JoinedGen = (s for s in joined if s[0]['source_type'] == 'navigation')
  events: JoinedGen = (s for s in joined if s[0]['source_type'] == 'event')

  if args.output_mode == 'aggregate':
    print(json.dumps({
        'navigation': correct_aggregates(navs, NAV_PARAMS),
        'event': correct_aggregates(events, EVENT_PARAMS),
    }, indent=2))

  elif args.output_mode == 'experimental-event-level':
    nav_corrected = generate_corrected_event_level(list(navs), NAV_PARAMS)
    event_corrected = generate_corrected_event_level(list(events), EVENT_PARAMS)
    combined = [{'source': s, 'reports': r}
                for s, r in nav_corrected + event_corrected]
    print(json.dumps(combined, indent=2))


if __name__ == '__main__':
  main()