57 lines
2.1 KiB
Python
57 lines
2.1 KiB
Python
from typing import Dict, Optional, Union
|
|
|
|
Fraction = float
|
|
Weight = float
|
|
Probability = float
|
|
Distribution = Union[
|
|
# For example, 1 means 100% reads.
|
|
int,
|
|
# For example, 0.25 means 25% reads.
|
|
float,
|
|
# For example, {0.25: 1, 0.8: 2} means 25% reads one third of the time and
|
|
# 80% reads two thirds of the time.
|
|
Dict[Fraction, Weight],
|
|
]
|
|
|
|
def canonicalize(d: Distribution) -> Dict[Fraction, Probability]:
|
|
if isinstance(d, int):
|
|
if d < 0 or d > 1:
|
|
raise ValueError('distribution must be in the range [0, 1]')
|
|
return {float(d): 1.}
|
|
elif isinstance(d, float):
|
|
if d < 0 or d > 1:
|
|
raise ValueError('distribution must be in the range [0, 1]')
|
|
return {d: 1.}
|
|
elif isinstance(d, dict):
|
|
if len(d) == 0:
|
|
raise ValueError('distribution cannot empty')
|
|
|
|
if any(weight < 0 for weight in d.values()):
|
|
raise ValueError('distribution cannot have negative weights')
|
|
|
|
total_weight = sum(d.values())
|
|
if total_weight == 0:
|
|
raise ValueError('distribution cannot have zero weight')
|
|
|
|
return {float(f): weight / total_weight
|
|
for (f, weight) in d.items()
|
|
if weight > 0}
|
|
else:
|
|
raise ValueError('distribution must be an int, a float, a Dict[float, '
|
|
'float] or a List[Tuple[float, float]]')
|
|
|
|
|
|
def canonicalize_rw(read_fraction: Optional[Distribution],
|
|
write_fraction: Optional[Distribution]) \
|
|
-> Dict[Fraction, Probability]:
|
|
if read_fraction is None and write_fraction is None:
|
|
raise ValueError('Either read_fraction or write_fraction must be given')
|
|
elif read_fraction is not None and write_fraction is not None:
|
|
raise ValueError('Only one of read_fraction or write_fraction can be '
|
|
'given')
|
|
elif read_fraction is not None:
|
|
return canonicalize(read_fraction)
|
|
else:
|
|
assert write_fraction is not None
|
|
return {1 - f: weight
|
|
for (f, weight) in canonicalize(write_fraction).items()}
|