# -*- coding: utf-8 -*-

#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You 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.

#

"""A workflow that uses a simple Monte Carlo method to estimate π.

The algorithm computes the fraction of points drawn uniformly within the unit

square that also fall in the quadrant of the unit circle that overlaps the

square. A simple area calculation shows that this fraction should be π/4, so

we multiply our counts ratio by four to estimate π.

"""

# pytype: skip-file

import argparse

import json

import logging

import random

from collections.abc import Iterable

from typing import Any

import apache_beam as beam

from apache_beam.io import WriteToText

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.options.pipeline_options import SetupOptions

@beam.typehints.with_output_types(tuple[int, int, int])

@beam.typehints.with_input_types(int)

def run_trials(runs):

"""Run trials and return a 3-tuple representing the results.

Args:

runs: Number of trial runs to be executed.

Returns:

A 3-tuple (total trials, inside trials, 0).

The final zero is needed solely to make sure that the combine_results function

has same type for inputs and outputs (a requirement for combiner functions).

"""

inside_runs = 0

for _ in range(runs):

x = random.uniform(0, 1)

y = random.uniform(0, 1)

inside_runs += 1 if x * x + y * y <= 1.0 else 0

return runs, inside_runs, 0

@beam.typehints.with_output_types(tuple[int, int, float])

@beam.typehints.with_input_types(Iterable[tuple[int, int, Any]])

def combine_results(results):

"""Combiner function to sum up trials and compute the estimate.

Args:

results: An iterable of 3-tuples (total trials, inside trials, ignored).

Returns:

A 3-tuple containing the sum of total trials, sum of inside trials, and

the probability computed from the two numbers.

"""

# TODO(silviuc): Do we guarantee that argument can be iterated repeatedly?

# Should document one way or the other.

total, inside = sum(r[0] for r in results), sum(r[1] for r in results)

return total, inside, 4 * float(inside) / total

class JsonCoder(object):

"""A JSON coder used to format the final result."""

def encode(self, x):

return json.dumps(x).encode('utf-8')

class EstimatePiTransform(beam.PTransform):

"""Runs 10M trials, and combine the results to estimate pi."""

def __init__(self, tries_per_work_item=100000):

self.tries_per_work_item = tries_per_work_item

def expand(self, pcoll):

# A hundred work items of a hundred thousand tries each.

return (

pcoll

| 'Initialize' >> beam.Create(

[self.tries_per_work_item] * 100).with_output_types(int)

| 'Run trials' >> beam.Map(run_trials)

| 'Sum' >> beam.CombineGlobally(combine_results).without_defaults())

def run(argv=None):

parser = argparse.ArgumentParser()

parser.add_argument(

'--output', required=True, help='Output file to write results to.')

known_args, pipeline_args = parser.parse_known_args(argv)

# We use the save_main_session option because one or more DoFn's in this

# workflow rely on global context (e.g., a module imported at module level).

pipeline_options = PipelineOptions(pipeline_args)

pipeline_options.view_as(SetupOptions).save_main_session = True

with beam.Pipeline(options=pipeline_options) as p:

( # pylint: disable=expression-not-assigned

p

| EstimatePiTransform()

| WriteToText(known_args.output, coder=JsonCoder()))

if __name__ == '__main__':

logging.getLogger().setLevel(logging.INFO)

run()