--fanout FLAG added to [weatherbench2]/scripts/compute_zonal_energy_spectrum

langmore · Weatherbench2 authors · commit aa86327fb239 · 2023-10-19T17:20:12.000-07:00
PiperOrigin-RevId: 575047105
diff --git a/scripts/compute_zonal_energy_spectrum.py b/scripts/compute_zonal_energy_spectrum.py
@@ -91,6 +91,11 @@
         ' average.'
     ),
 )
+FANOUT = flags.DEFINE_integer(
+    'fanout',
+    None,
+    help='Beam CombineFn fanout. Might be required for large dataset.',
+)
 
 RUNNER = flags.DEFINE_string('runner', None, 'beam.runners.Runner')
 
@@ -199,7 +204,7 @@ def main(argv: list[str]) -> None:
         )
         | xbeam.SplitVariables()
         | beam.MapTuple(_strip_offsets)
-        | xbeam.Mean(AVERAGING_DIMS.value)
+        | xbeam.Mean(AVERAGING_DIMS.value, fanout=FANOUT.value)
         | xbeam.ChunksToZarr(
             OUTPUT_PATH.value, template, output_chunks, num_threads=16
         )

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,11 @@`
`91`	`91`	`' average.'`
`92`	`92`	`),`
`93`	`93`	`)`
	`94`	`+FANOUT = flags.DEFINE_integer(`
	`95`	`+ 'fanout',`
	`96`	`+ None,`
	`97`	`+ help='Beam CombineFn fanout. Might be required for large dataset.',`
	`98`	`+)`
`94`	`99`
`95`	`100`	`RUNNER = flags.DEFINE_string('runner', None, 'beam.runners.Runner')`
`96`	`101`
`@@ -199,7 +204,7 @@ def main(argv: list[str]) -> None:`
`199`	`204`	`)`
`200`	`205`	`\| xbeam.SplitVariables()`
`201`	`206`	`\| beam.MapTuple(_strip_offsets)`
`202`		`- \| xbeam.Mean(AVERAGING_DIMS.value)`
	`207`	`+ \| xbeam.Mean(AVERAGING_DIMS.value, fanout=FANOUT.value)`
`203`	`208`	`\| xbeam.ChunksToZarr(`
`204`	`209`	`OUTPUT_PATH.value, template, output_chunks, num_threads=16`
`205`	`210`	`)`