Check for kwarg updates to track_matcher on reload

[sgossage]

Previously, if a matcher was created via a prior SimulationProperties.load_step() call, it would not be created again (intended) and could not have its kwargs updated (unintended).

This allows the kwargs to be updated upon reload. It also adds a separate verbosity flag for TrackMatcher that can be set to isolate track matching output, rather than all step output, the idea being to avoid unwanted screen clutter.

[maxbriel]

@sgossage What do you mean exactly with a reload?

[sgossage]

@sgossage What do you mean exactly with a reload?

For example, if you've done

my_ini_path = 'population_params.ini'
sim_prop = SimulationProperties.from_ini(my_ini_path, load_steps=True)

and subsequently want to reload a step like this, to change some setting:

# define the evolution step tuple for step_disrupted with record_matching = True this time
step_tup = (DisruptedStep, {'record_matching':True})

# reload the step with these new options turned on
sim_prop.load_a_step("step_disrupted", step_tup)

The code

def check_step(self, metallicity, RNG, step_name, step_tup, verbose=False):        
    ...
    matcher_key = (metallicity, step_name)        
    if "track_matcher" in step_func.DEFAULT_KWARGS:            
        matcher_needed = matcher_key not in self.track_matchers            
        if matcher_needed:                
            # create TrackMatcher if needed               
            step_kwargs, matcher_kwargs = TrackMatcher.separate_kwargs(step_kwargs)
    ...

would prevent matcher_kwargs being separated from step_kwargs (because this was only done if the TrackMatcher object didn't exist), which created errors by itself due to the step subsequently getting kwargs it does not expect and that should have gone to the TrackMatcher.

So now, kwargs are always separated, and TrackMatcher objects are updated if load_a_step has been given any new kwargs related to matching.

[maxbriel]

Somewhere it should be documented that updating the args, does not allow you to change the matching stuff.

On the creation of the TrackMatcher, we train given a specific list, for example list_for_matching_HMS. When reloading the kwargs, you can change this list list_for_matching_HMS, but it will still use the old trained values.

Instead wouldn't it be better to completely reload the TrackMatcher? Or is this too much work/code changes?

[sgossage]

Somewhere it should be documented that updating the args, does not allow you to change the matching stuff.

On the creation of the TrackMatcher, we train given a specific list, for example list_for_matching_HMS. When reloading the kwargs, you can change this list list_for_matching_HMS, but it will still use the old trained values.

Instead wouldn't it be better to completely reload the TrackMatcher? Or is this too much work/code changes?

I think that's a good catch. We can add something that recognizes if re-training is needed perhaps, and in that case reload the entire TrackMatcher, rather than just update kwargs.

I added something along these lines. It works for the evolution step you are updating, but does not enforce consistency across all evolution steps that have TrackMatchers, which may or may not be desired.

We now have something like

        ...
            if matcher_needed:
                # Always just create a new TrackMatcher if it does not exist
                ...
            else:
                track_matcher = self.track_matchers[matcher_key]
                retrain = False
                # subsequently check if any properties need to be updated,
                # in case reloading for example
                for k, v in matcher_kwargs.items():
                    # only care to update kwargs actually passed via load_step
                    if k in original_step_kwargs:
                        setattr(track_matcher, k, matcher_kwargs[k])
                        track_matcher.kwargs[k] = matcher_kwargs[k]
                        if k in track_matcher.TRAINING_TRIGGERS:
                            retrain = True
                if retrain:
                    updated_kwargs = track_matcher.kwargs
                    self.create_track_matcher(metallicity, step_name, updated_kwargs)

and in track_match.py:

            # these matcher kwargs will trigger retraining the TrackMatcher
            TRAINING_TRIGGERS = ["grid_Hrich", "grid_strippedHe", "path", "metallicity",
                                "list_for_matching_HMS", "list_for_matching_HeStar", 
                                "list_for_matching_postMS", "list_for_matching_postHeMS"]

[maxbriel]

The code block starting on line 520 always adds the metallicity to the step_kwargs. I think this means that this line is always triggered as TRAINING_TRIGGERS contains metallicity in it.

Does "always" retraining take a long time?

[sgossage]

Depending on the machine, I see recreation taking 10-40 sec or so. Most of the time is loading the grids, and retraining itself seems to take around 1-4 seconds.

Right now actually, if you change the metallicity, you populate SimulationProperties.track_matchers, which is a dictionary (keys are (metallicity, step_name))), with a new TrackMatcher and the other one reamains in memory. It would be more RAM efficient to delete the old one, so I've added a clean up step. I've also added a clean up step to check whether any cached grids (that were loaded at some metallicity) should be deleted. This is a niche scenario but should be safeguarded against now.

I've also organized most of the new code into a method called _check_track_matcher() that checks whether the current step under operation needs a TrackMatcher and takes care of either creating or updating it/retraining it and deleting old TrackMatchers or unneeded grids.

About metallicity: I've moved original_step_kwargs = step_kwargs.copy() to the top of the function call for check_step(), before the default metallicity is added to step_kwargs to make sure that it only holds values supplied in the original step_tup.

This line (now in _check_tack_matcher()) protects against retraining always firing though:

if k in original_step_kwargs:
    ...

Only kwargs that are passed in the original step_tup are considered for updates. I also added a guard so that only values that differ from what exists trigger an update via the do_update flag below:

            else:
                ...
                # subsequently check if any properties need to be updated,
                # in case reloading for example
                for k, v in matcher_kwargs.items():
                    # only care to update kwargs actually passed via load_step
                    if k in original_step_kwargs:
                        do_update = track_matcher.kwargs[k] != original_step_kwargs[k]
                        if do_update:
                            setattr(track_matcher, k, matcher_kwargs[k])
                            track_matcher.kwargs[k] = matcher_kwargs[k]
                            if k in track_matcher.TRAINING_TRIGGERS:
                                retrain = True

[maxbriel]

Thanks!