Skip to content

Commit f8623dc

Browse files
committed
MAINT: address Ruff LOG015 error
1 parent e151210 commit f8623dc

File tree

1 file changed

+10
-14
lines changed

1 file changed

+10
-14
lines changed

scripts/execute_jax_benchmark.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,8 @@
3333
format="%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s",
3434
datefmt="%Y-%m-%d %H:%M:%S",
3535
)
36-
LOGGER = logging.getLogger()
36+
_LOGGER = logging.getLogger()
3737
mute_jax_warnings()
38-
3938
yaml.add_representer(defaultdict, Representer.represent_dict)
4039

4140
BENCHMARK_CASES = [
@@ -78,7 +77,7 @@ def main() -> int:
7877
existing_benchmark = t.get(n)
7978
if existing_benchmark is not None:
8079
if all(len(v) == NUMBER_OF_RUNS for v in existing_benchmark.values()):
81-
logging.warning(f"Benchmark for {n:,} events already exists")
80+
_LOGGER.warning(f"Benchmark for {n:,} events already exists")
8281
progress_bar.update(NUMBER_OF_RUNS)
8382
continue
8483
t[n] = defaultdict(list)
@@ -118,33 +117,30 @@ def create_amplitude_model() -> AmplitudeModel:
118117
def prepare_functions(
119118
model: AmplitudeModel,
120119
) -> tuple[ParametrizedBackendFunction, PositionalArgumentFunction]:
121-
original_log_level = LOGGER.getEffectiveLevel()
122-
LOGGER.setLevel(logging.INFO)
123-
logging.info("Unfolding intensity expression")
120+
_LOGGER.info("Unfolding intensity expression")
124121
unfolded_intensity_expr = perform_cached_doit(model.full_expression)
125-
logging.info("Substituting parameters")
122+
_LOGGER.info("Substituting parameters")
126123
substituted_expr = unfolded_intensity_expr.xreplace(model.parameter_defaults)
127-
logging.info("Lambdifying full intensity expression")
124+
_LOGGER.info("Lambdifying full intensity expression")
128125
parametrized_func = create_parametrized_function(
129126
unfolded_intensity_expr,
130127
parameters=model.parameter_defaults,
131128
backend="jax",
132129
)
133-
logging.info("Lambdifying substituted intensity expression")
130+
_LOGGER.info("Lambdifying substituted intensity expression")
134131
substituted_func = create_function(substituted_expr, backend="jax")
135-
logging.info("Finished function lambdification")
136-
LOGGER.setLevel(original_log_level)
132+
_LOGGER.info("Finished function lambdification")
137133
return parametrized_func, substituted_func
138134

139135

140136
def generate_sample(
141137
model: AmplitudeModel, n_events: int, seed: int | None = None
142138
) -> DataSample:
143139
transformer = create_data_transformer(model)
144-
original_log_level = LOGGER.getEffectiveLevel()
145-
LOGGER.setLevel(logging.ERROR)
140+
original_log_level = _LOGGER.getEffectiveLevel()
141+
_LOGGER.setLevel(logging.ERROR)
146142
phsp_sample = generate_phasespace_sample(model.decay, n_events, seed)
147-
LOGGER.setLevel(original_log_level)
143+
_LOGGER.setLevel(original_log_level)
148144
return transformer(phsp_sample)
149145

150146

0 commit comments

Comments
 (0)