def test_determine_leading_hadron_bias(logging_mixin, bias_type, event_activity, expected_leading_hadron_bias_value, leading_hadron_bias_config, override_options_helper):
""" Test determination of the leading hadron bias. """
(config, selected_analysis_options) = override_options_helper(leading_hadron_bias_config)
# Add in the different selected options
if bias_type:
# Both options will lead here. Doing this with "track" doesn't change the values, but
# also doesn't hurt anything, so it's fine.
kwargs = dict(selected_analysis_options)
kwargs["leading_hadron_bias"] = params.LeadingHadronBiasType[bias_type]
selected_analysis_options = params.SelectedAnalysisOptions(**kwargs)
if event_activity:
kwargs = dict(selected_analysis_options)
kwargs["event_activity"] = params.EventActivity[event_activity]
selected_analysis_options = params.SelectedAnalysisOptions(**kwargs)
returned_options = analysis_config.determine_leading_hadron_bias(config = config, selected_analysis_options = selected_analysis_options)
# Check that we still got these right.
assert returned_options.collision_energy == selected_analysis_options.collision_energy
assert returned_options.collision_system == selected_analysis_options.collision_system
assert returned_options.event_activity == selected_analysis_options.event_activity
# Check the bias value and type
logger.debug(f"type(leading_hadron_bias): {type(returned_options.leading_hadron_bias)}")
assert returned_options.leading_hadron_bias.value == expected_leading_hadron_bias_value
assert returned_options.leading_hadron_bias.type == params.LeadingHadronBiasType[bias_type]
def test_validate_arguments(logging_mixin, args, expected):
""" Test argument validation. """
if expected is None:
expected = (params.CollisionEnergy.two_seven_six,
params.CollisionSystem.PbPb,
params.EventActivity.central,
params.LeadingHadronBiasType.track)
args = params.SelectedAnalysisOptions(*args)
args, _ = analysis_config.validate_arguments(args)
expected = params.SelectedAnalysisOptions(*expected)
assert args.collision_energy == expected.collision_energy
assert args.collision_system == expected.collision_system
assert args.event_activity == expected.event_activity
assert args.leading_hadron_bias == expected.leading_hadron_bias
def determine_leading_hadron_bias(
config: generic_config.DictLike,
selected_analysis_options: params.SelectedAnalysisOptions
) -> params.SelectedAnalysisOptions:
""" Determines the leading hadron bias based on the analysis options.
The determined leading hadron bias object is then created and stored in an updated selected
analysis options object which is returned.
Args:
config: Contains the dict-like analysis configuration. Note that it must already be
fully configured and overridden.
selected_analysis_options: Selected analysis options.
Returns:
Selected analysis options with the determined leading hadron bias object.
"""
override_options = generic_config.determine_override_options(
selected_options=selected_analysis_options.astuple(),
override_opts=config["leadingHadronBiasValues"],
set_of_possible_options=params.set_of_possible_options,
)
leading_hadron_bias_type = selected_analysis_options.leading_hadron_bias
leading_hadron_bias_value = override_options["value"]
# Help out mypy
assert isinstance(leading_hadron_bias_type, params.LeadingHadronBiasType)
# Namedtuple is immutable, so we need to return a new one with the proper parameters
return_options = dict(selected_analysis_options)
return_options["leading_hadron_bias"] = params.LeadingHadronBias(
type=leading_hadron_bias_type, value=leading_hadron_bias_value)
return params.SelectedAnalysisOptions(**return_options)
def test_JetHBase_object_construction(logging_mixin, leading_hadron_bias,
object_yaml_config,
override_options_helper,
check_JetHBase_object, mocker):
""" Test construction of the JetHBase object. """
object_config, task_name = object_yaml_config
(config, selected_analysis_options) = override_options_helper(
object_config, config_containing_override=object_config[task_name])
# Avoid os.makedirs actually making directories
mocker.patch("os.makedirs")
config_filename = "configFilename.yaml"
task_config = config[task_name]
reaction_plane_orientation = params.ReactionPlaneOrientation.inclusive
config_base = analysis_objects.JetHBase(
task_name=task_name,
config_filename=config_filename,
config=config,
task_config=task_config,
collision_energy=selected_analysis_options.collision_energy,
collision_system=selected_analysis_options.collision_system,
event_activity=selected_analysis_options.event_activity,
leading_hadron_bias=selected_analysis_options.leading_hadron_bias,
reaction_plane_orientation=reaction_plane_orientation,
)
# We need values to compare against. However, namedtuples are immutable,
# so we have to create a new one with the proper value.
temp_selected_options = dict(selected_analysis_options)
temp_selected_options["leading_hadron_bias"] = leading_hadron_bias
selected_analysis_options = params.SelectedAnalysisOptions(
**temp_selected_options)
# Only need for the case of LeadingHadronBiasType!
if isinstance(leading_hadron_bias, params.LeadingHadronBiasType):
selected_analysis_options = analysis_config.determine_leading_hadron_bias(
config, selected_analysis_options)
# Assertions are performed in this function
res = check_JetHBase_object(
obj=config_base,
config=config,
selected_analysis_options=selected_analysis_options,
reaction_plane_orientation=reaction_plane_orientation)
assert res is True
# Just to be safe
mocker.stopall()
def func(config, selected_options=None, config_containing_override=None):
""" Helper function to override the configuration.
It can print the configuration before and after overridding the options if enabled.
Note:
If selected_options is not specified, it defaults to (2.76, "PbPb", "central", "track")
Args:
config (CommentedMap): dict-like object containing the configuration to be overridden.
selected_options (params.SelectedAnalysisOptions): The options selected for this analysis, in
the order defined used with analysis_config.override_options() and in the configuration file.
config_containing_override (CommentedMap): dict-like object containing the override options.
Returns:
tuple: (dict-like CommentedMap object containing the overridden configuration, selected analysis
options used with the config)
"""
# Import modules here so we can delay it until they are actually needed.
from jet_hadron.base import analysis_config
from jet_hadron.base import params
import logging
logger = logging.getLogger(__name__)
if selected_options is None:
selected_options = params.SelectedAnalysisOptions(
collision_energy=params.CollisionEnergy.two_seven_six,
collision_system=params.CollisionSystem.PbPb,
event_activity=params.EventActivity.central,
leading_hadron_bias=params.LeadingHadronBiasType.track)
yaml = ruamel.yaml.YAML()
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Before override:")
logger.debug(log_yaml_dump(yaml, config))
config = analysis_config.override_options(
config=config,
selected_options=selected_options,
config_containing_override=config_containing_override)
if logger.isEnabledFor(logging.DEBUG):
logger.debug("After override:")
logger.debug(log_yaml_dump(yaml, config))
return (config, selected_options)
def leading_hadron_bias(self) -> params.LeadingHadronBias:
# Only calculate the value if we haven't already used it.
if self._leading_hadron_bias is None:
# Load this module only if necessary. I'm not moving this function because it makes dependences much messier.
from jet_hadron.base import analysis_config
self._leading_hadron_bias = analysis_config.determine_leading_hadron_bias( # type: ignore
config = self.config,
selected_analysis_options = params.SelectedAnalysisOptions(
collision_energy = self.collision_energy,
collision_system = self.collision_system,
event_activity = self.event_activity,
leading_hadron_bias = self._leading_hadron_bias_type)
).leading_hadron_bias
# Help out mypy.
assert self._leading_hadron_bias is not None
return self._leading_hadron_bias
assert qvector.value_range == expected["range"]
# Integration tests
@pytest.mark.parametrize("obj", [
params.SelectedRange(min=-5, max=15),
params.ReactionPlaneBinInformation(bin=1, center=0, width=np.pi / 6),
params.CollisionEnergy.five_zero_two,
params.CollisionSystem.embedPythia,
params.EventActivity.semi_central,
params.LeadingHadronBiasType.track,
params.LeadingHadronBias(type=params.LeadingHadronBiasType.cluster,
value=6.0),
params.SelectedAnalysisOptions(
collision_system=params.CollisionSystem.PbPb,
collision_energy=params.CollisionEnergy.two_seven_six,
event_activity=params.EventActivity.central,
leading_hadron_bias=params.LeadingHadronBiasType.track,
),
params.ReactionPlaneOrientation.out_of_plane,
params.QVector.bottom10,
],
ids=[
"SelectedRange", "ReactionPlaneBinInformation",
"CollisionEnergy", "CollisionSystem",
"EventActivity", "LeadingHadronBiasType",
"LeadingHadronBias", "SelectedAnalysisOptions",
"ReactionPlaneOrientation", "QVector"
])
def test_yaml_round_trip(logging_mixin, dump_to_string_and_retrieve, obj):
""" Integrations tests for writing to and reading from YAML. """
# Setup
def determine_selected_options_from_kwargs(
args: Optional[List[Any]] = None,
description: str = "Jet-hadron {task_name}.",
add_options_function: Optional[Callable[[argparse.ArgumentParser],
Any]] = None,
**kwargs: str
) -> Tuple[str, params.SelectedAnalysisOptions, argparse.Namespace]:
""" Determine the selected analysis options from the command line arguments.
Defaults are equivalent to None or False so values can be added in the validation
function if argument values are not specified.
Args:
args (list): Arguments to parse. Default: None (which will then use sys.argv)
description (str): Help description for arguments
add_options_function (func): Function which takes the ArgumentParser() object, adds
arguments, and returns the object.
kwargs (dict): Additional arguments to format the help description. Often contains ``task_name``
to specify the task name.
Returns:
tuple: (config_filename, energy, collision_system, event_activity, bias_type, argparse.Namespace).
The args are return for handling custom arguments added with add_options_function.
"""
# Make sure there is always a task name
if "task_name" not in kwargs:
kwargs["task_name"] = "analysis"
# Setup parser
parser = argparse.ArgumentParser(description=description.format(**kwargs))
# General options
parser.add_argument("-c",
"--configFilename",
metavar="configFilename",
type=str,
default="config/analysisConfig.yaml",
help="Path to config filename")
parser.add_argument("-e",
"--energy",
metavar="energy",
type=float,
default=0.0,
help="Collision energy")
parser.add_argument("-s",
"--collisionSystem",
metavar="collisionSystem",
type=str,
default="",
help="Collision system")
parser.add_argument("-a",
"--eventActivity",
metavar="eventActivity",
type=str,
default="",
help="Event activity")
parser.add_argument("-b",
"--biasType",
metavar="biasType",
type=str,
default="",
help="Leading hadron bias type")
# Extension for additional arguments
if add_options_function:
args = add_options_function(parser)
# Parse arguments
parsed_args = parser.parse_args(args)
# Even though we will need to create a new selected analysis options tuple, we store the
# return values in one for convenience.
selected_analysis_options = params.SelectedAnalysisOptions(
collision_energy=parsed_args.energy,
collision_system=parsed_args.collisionSystem,
event_activity=parsed_args.eventActivity,
leading_hadron_bias=parsed_args.biasType)
return (parsed_args.configFilename, selected_analysis_options, parsed_args)
def validate_arguments(
selected_args: params.SelectedAnalysisOptions,
validate_extra_args_func: Any = None
) -> Tuple[params.SelectedAnalysisOptions, Dict[str, Any]]:
""" Validate arguments passed to the analysis task. Converts str and float types to enumerations.
Note:
If the selections are not specified, it will define to 2.76 TeV central PbPb collisions with a
track bias!
Args:
selected_args: Selected analysis options from args or otherwise.
validate_extra_args_func (func): Function to validate additional args that were added using
``add_options_function()``. It should be a closure with the args returned from initial parsing
and return a dict containing the validated args.
Returns:
tuple: (validated_selected_options, additional_validated_args)
"""
# Validate the given arguments.
# The general strategy is as follows:
# Input:
# - If the value is None, use the default.
# - If the value is given, then use the given value.
# Enum object creation:
# - Check if the input is already of the enum type. If so, use it.
# - If not, initialize the enum value using the given value.
# Energy. Default: 2.76
energy = selected_args.collision_energy if selected_args.collision_energy else 2.76
# Retrieves the enum by value
energy = energy if type(
energy) is params.CollisionEnergy else params.CollisionEnergy(energy)
# Help out mypy...
assert isinstance(energy, params.CollisionEnergy)
# Collision system. Default: PbPb
collision_system = selected_args.collision_system if selected_args.collision_system else "PbPb"
collision_system = collision_system if type(
collision_system
) is params.CollisionSystem else params.CollisionSystem[
collision_system] # type: ignore
# Help out mypy...
assert isinstance(collision_system, params.CollisionSystem)
# Event activity. Default: central
event_activity = selected_args.event_activity if selected_args.event_activity else "central"
event_activity = event_activity if type(
event_activity) is params.EventActivity else params.EventActivity[
event_activity] # type: ignore
# Help out mypy...
assert isinstance(event_activity, params.EventActivity)
# Leading hadron bias type. Default: track
leading_hadron_bias_type = selected_args.leading_hadron_bias if selected_args.leading_hadron_bias else "track"
leading_hadron_bias_type = leading_hadron_bias_type if type(
leading_hadron_bias_type
) is params.LeadingHadronBiasType else params.LeadingHadronBiasType[
leading_hadron_bias_type] # type: ignore
# Help out mypy...
assert isinstance(leading_hadron_bias_type,
(params.LeadingHadronBias, params.LeadingHadronBiasType))
# Handle additional arguments
additional_validated_args: Dict[str, Any] = {}
if validate_extra_args_func:
additional_validated_args.update(validate_extra_args_func())
selected_analysis_options = params.SelectedAnalysisOptions(
collision_energy=energy,
collision_system=collision_system,
event_activity=event_activity,
leading_hadron_bias=leading_hadron_bias_type,
)
return (selected_analysis_options, additional_validated_args)