def sweep_from_proto(msg: run_context_pb2.Sweep) -> sweeps.Sweep:
"""Creates a Sweep from a v2 protobuf message."""
which = msg.WhichOneof('sweep')
if which is None:
return sweeps.UnitSweep
elif which == 'sweep_function':
factors = [sweep_from_proto(m) for m in msg.sweep_function.sweeps]
func_type = msg.sweep_function.function_type
if func_type == run_context_pb2.SweepFunction.PRODUCT:
return sweeps.Product(*factors)
elif func_type == run_context_pb2.SweepFunction.ZIP:
return sweeps.Zip(*factors)
else:
raise ValueError(
'invalid sweep function type: {}'.format(func_type))
elif which == 'single_sweep':
key = msg.single_sweep.parameter_key
if msg.single_sweep.WhichOneof('sweep') == 'linspace':
return sweeps.Linspace(
key=key,
start=msg.single_sweep.linspace.first_point,
stop=msg.single_sweep.linspace.last_point,
length=msg.single_sweep.linspace.num_points,
)
elif msg.single_sweep.WhichOneof('sweep') == 'points':
return sweeps.Points(key=key,
points=msg.single_sweep.points.points)
else:
raise ValueError('single sweep type not set: {}'.format(msg))
else:
# coverage: ignore
raise ValueError('sweep type not set: {}'.format(msg))
def _sweep_from_param_sweep_zip_proto_dict(
param_sweep_zip: Dict) -> sweeps.Sweep:
if 'sweeps' in param_sweep_zip:
return sweeps.Zip(*[
_sweep_from_single_param_sweep_proto_dict(sweep)
for sweep in param_sweep_zip['sweeps']
])
return sweeps.UnitSweep
def _sweep_from_param_sweep_zip_proto(
param_sweep_zip: params_pb2.ZipSweep) -> sweeps.Sweep:
if len(param_sweep_zip.sweeps) > 0:
return sweeps.Zip(*[
_sweep_from_single_param_sweep_proto(sweep)
for sweep in param_sweep_zip.sweeps
])
return sweeps.UnitSweep
def _to_zip_product(sweep: sweeps.Sweep) -> sweeps.Product:
"""Converts sweep to a product of zips of single sweeps, if possible."""
if not isinstance(sweep, sweeps.Product):
sweep = sweeps.Product(sweep)
if not all(isinstance(f, sweeps.Zip) for f in sweep.factors):
factors = [f if isinstance(f, sweeps.Zip) else sweeps.Zip(f) for f in sweep.factors]
sweep = sweeps.Product(*factors)
for factor in sweep.factors:
for term in cast(sweeps.Zip, factor).sweeps:
if not isinstance(term, sweeps.SingleSweep):
raise ValueError(f'cannot convert to zip-product form: {sweep}')
return sweep