Faster parent population trials lookup in NSGA-II sampler

[c-bata]

Summary

Follow-up #179 and #222.

• Resolve persisted NSGA-II parent trial IDs through get_cached_trial.
• Avoid rebuilding the completed trial_id -> trial number map by scanning all trials on every sample.

Motivation

NSGA-II parent populations are persisted as trial IDs. On main, restoring a parent population rebuilds a trial_id -> trial number map by scanning all completed trials for every sample_joint call. This makes the cost grow with the total number of trials, even though the parent population contains only population_size trials.

Benchmark

I took the same benchmark with #237 (comment). Here are the benchmark results (n_trials=10000, n_params=40):

branch	elapsed
main branch	128.034s
PR #179	7.878s
PR #237	123.207s
This PR	7.630s

[c-bata]

I also executed the following benchmark to confirm that this PR does not change the behavior.

import time
import rustuna

n_trials = 10000


def run_optimize():
    def objective(trial: rustuna.Trial) -> tuple[float, float]:
        x = trial.suggest_float("x", -15, 30)
        y = trial.suggest_float("y", -15, 30)

        v0 = 4 * x**2 + 4 * y**2
        v1 = (x - 5) ** 2 + (y - 5) ** 2
        trial.set_constraints({"c0": 1000 - v0})
        return v0, v1

    directions = ["minimize", "minimize"]
    sampler = rustuna.samplers.NSGAIISampler(seed=1)
    study = rustuna.create_study(sampler=sampler, directions=directions)
    study.optimize(objective, n_trials=n_trials)
    return study


def main():
    for _ in range(3):
        start = time.time()
        study = run_optimize()
        elapsed_rustuna = time.time() - start
        assert len(study.get_trials(states=[rustuna.trial.TrialState.COMPLETE])) == n_trials
        print(f"Rustuna\telapsed={elapsed_rustuna:.3f}\tbest_trials[0].values={study.best_trials[0].values}")

if __name__ == "__main__":
    main()

main branch:

% python ./bench_ngaii.py
Rustuna elapsed=1.286   best_trials[0].values=[9.837239422697502, 30.588169844860616]
Rustuna elapsed=1.249   best_trials[0].values=[9.837239422697502, 30.588169844860616]
Rustuna elapsed=1.226   best_trials[0].values=[9.837239422697502, 30.588169844860616]

This PR:

% python ./bench_ngaii.py
Rustuna elapsed=0.099   best_trials[0].values=[9.837239422697502, 30.588169844860616]
Rustuna elapsed=0.079   best_trials[0].values=[9.837239422697502, 30.588169844860616]
Rustuna elapsed=0.080   best_trials[0].values=[9.837239422697502, 30.588169844860616]

[c-bata]

@Alnusjaponica Could you review this PR?

[Alnusjaponica]

Apologies for the delayed reply. I've added a few comments about the behavioral changes. Please take a look, and let me know if anything is unclear.

[c-bata]

@Alnusjaponica I applied your suggestions. PTAL.

[Alnusjaponica]

LGTM. Thanks for the update.

[Alnusjaponica]

ref: