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)