Skip to content

Commit

Permalink
removed 'round' variable name clash with builtin
Browse files Browse the repository at this point in the history
  • Loading branch information
tgalvin committed Oct 11, 2024
1 parent 17e80fe commit ae7c2af
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 42 deletions.
49 changes: 28 additions & 21 deletions flint/configuration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,7 @@ def get_image_options_from_yaml(
def get_options_from_strategy(
strategy: Union[Strategy, None, Path],
mode: str = "wsclean",
round: Union[str, int] = "initial",
round_info: Union[str, int] = "initial",
max_round_override: bool = True,
operation: Optional[str] = None,
) -> Dict[Any, Any]:
Expand All @@ -241,7 +241,7 @@ def get_options_from_strategy(
Args:
strategy (Union[Strategy,None,Path]): A loaded instance of a strategy file. If `None` is provided then an empty dictionary is returned. If `Path` attempt to load the strategy file.
mode (str, optional): Which set of options to load. Typical values are `wsclean`, `gaincal` and `masking`. Defaults to "wsclean".
round (Union[str, int], optional): Which round to load options for. May be `initial` or an `int` (which indicated a self-calibration round). Defaults to "initial".
round_info (Union[str, int], optional): Which round to load options for. May be `initial` or an `int` (which indicated a self-calibration round). Defaults to "initial".
max_round_override (bool, optional): Check whether an integer round number is recorded. If it is higher than the largest self-cal round specified, set it to the last self-cal round. If False this is not performed. Defaults to True.
operation (Optional[str], optional): Get options related to a specific operation. Defaults to None.
Expand All @@ -262,13 +262,17 @@ def get_options_from_strategy(
assert isinstance(
strategy, (Strategy, dict)
), f"Unknown input strategy type {type(strategy)}"
assert round == "initial" or isinstance(
round, int
), f"{round=} not a known value or type. "
assert round_info == "initial" or isinstance(
round_info, int
), f"{round_info=} not a known value or type. "

# Override the round if requested
if isinstance(round, int) and max_round_override and "selfcal" in strategy.keys():
round = min(round, max(strategy["selfcal"].keys()))
if (
isinstance(round_info, int)
and max_round_override
and "selfcal" in strategy.keys()
):
round_info = min(round_info, max(strategy["selfcal"].keys()))

# step one, get the defaults
options = dict(**strategy["defaults"][mode]) if mode in strategy["defaults"] else {}
Expand All @@ -286,16 +290,19 @@ def get_options_from_strategy(
)
if mode in strategy[operation]:
update_options = dict(**strategy[operation][mode])
elif round == "initial":
elif round_info == "initial":
# separate function to avoid a missing mode from raising value error
if mode in strategy["initial"]:
update_options = dict(**strategy["initial"][mode])
elif isinstance(round, int):
elif isinstance(round_info, int):
# separate function to avoid a missing mode from raising value error
if round in strategy["selfcal"] and mode in strategy["selfcal"][round]:
update_options = dict(**strategy["selfcal"][round][mode])
if (
round_info in strategy["selfcal"]
and mode in strategy["selfcal"][round_info]
):
update_options = dict(**strategy["selfcal"][round_info][mode])
else:
raise ValueError(f"{round=} not recognised.")
raise ValueError(f"{round_info=} not recognised.")

if update_options:
logger.debug(f"Updating options with {update_options=}")
Expand Down Expand Up @@ -344,7 +351,7 @@ def wrapper(
*args,
strategy: Union[Strategy, None, Path] = None,
mode: str = "wsclean",
round: Union[str, int] = "initial",
round_info: Union[str, int] = "initial",
max_round_override: bool = True,
operation: Optional[str] = None,
**kwargs,
Expand All @@ -357,7 +364,7 @@ def wrapper(
update_options = get_options_from_strategy(
strategy=strategy,
mode=mode,
round=round,
round_info=round_info,
max_round_override=max_round_override,
operation=operation,
)
Expand Down Expand Up @@ -408,7 +415,7 @@ def verify_configuration(input_strategy: Strategy, raise_on_error: bool = True)
else:
for key in input_strategy["initial"].keys():
options = get_options_from_strategy(
strategy=input_strategy, mode=key, round="initial"
strategy=input_strategy, mode=key, round_info="initial"
)
try:
_ = MODE_OPTIONS_MAPPING[key](**options)
Expand All @@ -423,16 +430,16 @@ def verify_configuration(input_strategy: Strategy, raise_on_error: bool = True)
if not all([isinstance(i, int) for i in round_keys]):
errors.append("The keys into the self-calibration should be ints. ")

for round in round_keys:
for mode in input_strategy["selfcal"][round]:
for round_info in round_keys:
for mode in input_strategy["selfcal"][round_info]:
options = get_options_from_strategy(
strategy=input_strategy, mode=mode, round=round
strategy=input_strategy, mode=mode, round_info=round_info
)
try:
_ = MODE_OPTIONS_MAPPING[mode](**options)
except TypeError as typeerror:
errors.append(
f"{mode=} mode in {round=} incorrectly formed. {typeerror} "
f"{mode=} mode in {round_info=} incorrectly formed. {typeerror} "
)

for operation in KNOWN_OPERATIONS:
Expand Down Expand Up @@ -534,8 +541,8 @@ def create_default_yaml(
if selfcal_rounds:
logger.info(f"Creating {selfcal_rounds} self-calibration rounds. ")
selfcal: Dict[int, Any] = {}
for round in range(1, selfcal_rounds + 1):
selfcal[round] = {
for selfcal_round in range(1, selfcal_rounds + 1):
selfcal[selfcal_round] = {
"wsclean": {},
"gaincal": {},
"masking": {},
Expand Down
10 changes: 5 additions & 5 deletions flint/prefect/flows/continuum_pipeline.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -281,7 +281,7 @@ def process_science_fields(
# The call into potato peel task has two potential update option keywords.
# So for the moment we will not use the task decorated version.
potato_wsclean_init = get_options_from_strategy(
strategy=strategy, mode="wsclean", round="initial"
strategy=strategy, mode="wsclean", round_info="initial"
)
preprocess_science_mss = task_potato_peel.map(
ms=preprocess_science_mss,
Expand All @@ -295,7 +295,7 @@ def process_science_fields(
wsclean_container=field_options.wsclean_container,
strategy=unmapped(strategy),
mode="wsclean",
round="initial",
round_info="initial",
)
# TODO: This should be waited!
beam_summaries = task_create_beam_summary.map(
Expand Down Expand Up @@ -368,7 +368,7 @@ def process_science_fields(
casa_container=field_options.casa_container,
strategy=unmapped(strategy),
mode="gaincal",
round=current_round,
round_info=current_round,
wait_for=[
field_summary
], # To make sure field summary is created with unzipped MSs
Expand Down Expand Up @@ -399,7 +399,7 @@ def process_science_fields(
image_products=beam_aegean_outputs,
strategy=unmapped(strategy),
mode="masking",
round=current_round,
round_info=current_round,
)

wsclean_cmds = task_wsclean_imager.map(
Expand All @@ -408,7 +408,7 @@ def process_science_fields(
fits_mask=fits_beam_masks,
strategy=unmapped(strategy),
mode="wsclean",
round=current_round,
round_info=current_round,
)
archive_wait_for.extend(wsclean_cmds)

Expand Down
38 changes: 22 additions & 16 deletions tests/test_configuration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,19 +193,19 @@ def test_get_options(strategy):
# test to make sure we can generate a default strategy (see pytest fixture)
# read it backinto a dict and then access some attributes
wsclean = get_options_from_strategy(
strategy=strategy, mode="wsclean", round="initial"
strategy=strategy, mode="wsclean", round_info="initial"
)
assert isinstance(wsclean, dict)
# example options
assert wsclean["data_column"] == "CORRECTED_DATA"

wsclean = get_options_from_strategy(strategy=strategy, mode="wsclean", round=1)
wsclean = get_options_from_strategy(strategy=strategy, mode="wsclean", round_info=1)
assert isinstance(wsclean, dict)
# example options
assert wsclean["data_column"] == "CORRECTED_DATA"

archive = get_options_from_strategy(
strategy=strategy, mode="archive", round="initial"
strategy=strategy, mode="archive", round_info="initial"
)
assert isinstance(archive, dict)
assert len(archive) > 0
Expand All @@ -218,14 +218,14 @@ def test_get_options(strategy):
def test_get_mask_options(package_strategy):
"""Basic test to prove masking operation behaves well"""
masking = get_options_from_strategy(
strategy=package_strategy, mode="masking", round="initial"
strategy=package_strategy, mode="masking", round_info="initial"
)

assert isinstance(masking, dict)
assert masking["flood_fill_positive_seed_clip"] == 4.5

masking2 = get_options_from_strategy(
strategy=package_strategy, mode="masking", round=1
strategy=package_strategy, mode="masking", round_info=1
)

print(strategy)
Expand All @@ -245,14 +245,14 @@ def test_get_mask_options_from_path(package_strategy_path):
package_strategy = Path(package_strategy_path)

masking = get_options_from_strategy(
strategy=package_strategy, mode="masking", round="initial"
strategy=package_strategy, mode="masking", round_info="initial"
)

assert isinstance(masking, dict)
assert masking["flood_fill_positive_seed_clip"] == 4.5

masking2 = get_options_from_strategy(
strategy=package_strategy, mode="masking", round=1
strategy=package_strategy, mode="masking", round_info=1
)

assert isinstance(masking2, dict)
Expand Down Expand Up @@ -294,7 +294,7 @@ def test_wrapper_function_val(package_strategy_path):
val=in_val,
strategy=Path(package_strategy_path),
mode="masking",
round="initial",
round_info="initial",
)

assert val == in_val
Expand All @@ -319,7 +319,7 @@ def test_wrapper_function(package_strategy_path):
assert val == "Jack"

val, update_options = print_return_values(
strategy=Path(package_strategy_path), mode="masking", round="initial"
strategy=Path(package_strategy_path), mode="masking", round_info="initial"
)

assert val == "Jack"
Expand All @@ -339,18 +339,20 @@ def test_get_modes(package_strategy):
strategy = package_strategy

for mode in ("wsclean", "gaincal", "masking"):
test = get_options_from_strategy(strategy=strategy, mode=mode, round=1)
test = get_options_from_strategy(strategy=strategy, mode=mode, round_info=1)
assert isinstance(test, dict)
assert len(test.keys()) > 0


def test_bad_round(package_strategy):
# make sure incorrect round is handled properly
with pytest.raises(AssertionError):
_ = get_options_from_strategy(strategy=package_strategy, round="doesnotexists")
_ = get_options_from_strategy(
strategy=package_strategy, round_info="doesnotexists"
)

with pytest.raises(AssertionError):
_ = get_options_from_strategy(strategy=package_strategy, round=1.23456)
_ = get_options_from_strategy(strategy=package_strategy, round_info=1.23456)


def test_empty_strategy_empty_options():
Expand All @@ -366,7 +368,7 @@ def test_max_round_override(package_strategy):
# round is sound
strategy = package_strategy

opts = get_options_from_strategy(strategy=strategy, round=9999)
opts = get_options_from_strategy(strategy=strategy, round_info=9999)
assert isinstance(opts, dict)
assert opts["data_column"] == "TheLastRoundIs3"

Expand Down Expand Up @@ -395,15 +397,19 @@ def test_updated_get_options(package_strategy):
assert isinstance(strategy, Strategy)

wsclean_init = get_options_from_strategy(
strategy=strategy, mode="wsclean", round="initial"
strategy=strategy, mode="wsclean", round_info="initial"
)
assert wsclean_init["data_column"] == "CORRECTED_DATA"

wsclean_1 = get_options_from_strategy(strategy=strategy, mode="wsclean", round=1)
wsclean_1 = get_options_from_strategy(
strategy=strategy, mode="wsclean", round_info=1
)
assert wsclean_init["data_column"] == "CORRECTED_DATA"
assert wsclean_1["data_column"] == "EXAMPLE"

wsclean_2 = get_options_from_strategy(strategy=strategy, mode="wsclean", round=2)
wsclean_2 = get_options_from_strategy(
strategy=strategy, mode="wsclean", round_info=2
)
assert wsclean_2["multiscale"] is False
assert wsclean_2["data_column"] == "CORRECTED_DATA"

Expand Down

0 comments on commit ae7c2af

Please sign in to comment.