eckity-dts 0.1.0__tar.gz

eckity_dts-0.1.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Eliad Shem-Tov, Ron Edri, Achiya Elyasaf and the EC-KitY project
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

eckity_dts-0.1.0/PKG-INFO

@@ -0,0 +1,167 @@
+Metadata-Version: 2.4
+Name: eckity-dts
+Version: 0.1.0
+Summary: Deep Tournament Selection operator for EC-KitY genetic algorithms
+Author: Eliad Shem-Tov, Ron Edri, Achiya Elyasaf
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/EC-KitY/DTS-for-GA
+Project-URL: Repository, https://github.com/EC-KitY/DTS-for-GA
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: eckity~=0.4.1
+Requires-Dist: numpy>=2.0.2
+Requires-Dist: overrides>=7.7.0
+Requires-Dist: scipy>=1.13.0
+Requires-Dist: torch>=2.7.1
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+Requires-Dist: twine>=5.0; extra == "dev"
+Dynamic: license-file
+
+# Deep Tournament Selection for EC-KitY
+
+`eckity-dts` provides Deep Tournament Selection (DTS), a learned selection operator for genetic algorithms built on [EC-KitY](https://github.com/EC-KitY/EC-KitY).
+
+DTS uses a Transformer encoder and a self-attention pointer network trained online with REINFORCE. It was introduced in **“Deep Tournament Selection for Genetic Algorithms”** by Eliad Shem-Tov, Ron Edri, and Achiya Elyasaf. The paper has not yet been published; a formal citation will be added when available.
+
+## Installation
+
+```bash
+pip install eckity-dts
+```
+
+## Public API
+
+```python
+from eckity_dts import CachingEvaluator, DeepTournamentSelection, DTSPolicy
+```
+
+## Constructing DTS
+
+The policy combines a population encoder with a pointer network:
+
+```python
+from eckity_dts import DeepTournamentSelection, DTSPolicy
+from deep_tournament_selection.selection.population_to_vec_transformer import (
+    PopulationToVecTransformer,
+)
+from deep_tournament_selection.selection.self_attention_pointer import (
+    SelfAttentionPointer,
+)
+
+population_size = 100
+vocab_size = 2  # maximum gene value + 1
+
+encoder = PopulationToVecTransformer(
+    vocab_size=vocab_size,
+    emb_dim=32,
+    latent_dim=32,
+    n_heads=4,
+    n_layers=2,
+    dim_feedforward=256,
+    max_pointers=population_size,
+)
+
+pointer = SelfAttentionPointer(
+    pointer_len=population_size,
+    d_model=32,
+)
+
+policy = DTSPolicy(
+    pop_to_vec_transformer=encoder,
+    pointer_transformer=pointer,
+    device="cpu",
+    train_every_n_gens=10,
+    learning_rate=2e-3,
+    final_lr=1e-3,
+    epsilon_greedy=1.0,
+    epsilon_greedy_decay=0.999,
+    min_epsilon=0.2,
+)
+
+dts = DeepTournamentSelection(policy, higher_is_better=True)
+```
+
+Use it as the EC-KitY selection method:
+
+```python
+selection_methods=[(dts, 1)]
+```
+
+Important parameters:
+
+- `population_size` determines the rank-embedding and pointer-table capacities.
+- `vocab_size` is the maximum integer gene value plus one.
+- `train_every_n_gens` controls how often accumulated trajectories train the policy.
+- `epsilon_greedy` and its decay control teacher-forced tournament selection versus learned selection.
+- `custom_reward_function`, when supplied, receives current fitness, previous fitness, and population arrays.
+- `device` can be `"cpu"` or `"cuda"` when a compatible PyTorch installation is available.
+
+## Fitness caching
+
+EC-KitY may evaluate unchanged vectors again across generations. `CachingEvaluator` avoids recomputing fitness for vectors it has already seen:
+
+```python
+from eckity_dts import CachingEvaluator
+
+evaluator = CachingEvaluator(MyEvaluator())
+print(evaluator.cache_stats())
+```
+
+## Compatibility
+
+- Python 3.9 or newer
+- EC-KitY 0.4.x
+- NumPy 2.0.2 or newer
+- SciPy 1.13.0 or newer
+- PyTorch 2.7.1 or newer
+- overrides 7.7.0 or newer
+
+These bounds are compatible with `eckity-dnc`, `eckity-bert-ga`, and `eckity-bert-gp`. None of the operator packages depends directly on another operator package.
+
+## Research repository
+
+The repository contains the full paper experiments for Graph Coloring, Set Cover, and TSP, together with benchmark instances, notebook-style runners, and result figures. These research resources are not included in the `eckity-dts` wheel.
+
+For repository development:
+
+```bash
+uv sync --extra dev --resolution lowest-direct
+uv run pytest
+```
+
+Experiment entry points remain available from a source checkout:
+
+```bash
+python -m deep_tournament_selection.experiments.graph_coloring --instance queen8_12.col.txt --generations 200
+python -m deep_tournament_selection.experiments.set_cover --instance scp41.txt --generations 200
+python -m deep_tournament_selection.experiments.tsp --instance att48.tsp --generations 200
+python run_experiments.py --selection both --runs 3 --generations 500
+```
+
+Results are written under `runs/`. Paper figures are stored under `figures/`, and the architecture diagram is under `images/`.
+
+## Development and release
+
+```bash
+uv run pytest
+uv run ruff check .
+uv build
+```
+
+Release preparation and manual PyPI upload commands are documented in [`RELEASING.md`](RELEASING.md).
+
+## License
+
+This project is licensed under the BSD 3-Clause License. See [`LICENSE`](LICENSE).

eckity_dts-0.1.0/README.md

@@ -0,0 +1,136 @@
+# Deep Tournament Selection for EC-KitY
+
+`eckity-dts` provides Deep Tournament Selection (DTS), a learned selection operator for genetic algorithms built on [EC-KitY](https://github.com/EC-KitY/EC-KitY).
+
+DTS uses a Transformer encoder and a self-attention pointer network trained online with REINFORCE. It was introduced in **“Deep Tournament Selection for Genetic Algorithms”** by Eliad Shem-Tov, Ron Edri, and Achiya Elyasaf. The paper has not yet been published; a formal citation will be added when available.
+
+## Installation
+
+```bash
+pip install eckity-dts
+```
+
+## Public API
+
+```python
+from eckity_dts import CachingEvaluator, DeepTournamentSelection, DTSPolicy
+```
+
+## Constructing DTS
+
+The policy combines a population encoder with a pointer network:
+
+```python
+from eckity_dts import DeepTournamentSelection, DTSPolicy
+from deep_tournament_selection.selection.population_to_vec_transformer import (
+    PopulationToVecTransformer,
+)
+from deep_tournament_selection.selection.self_attention_pointer import (
+    SelfAttentionPointer,
+)
+
+population_size = 100
+vocab_size = 2  # maximum gene value + 1
+
+encoder = PopulationToVecTransformer(
+    vocab_size=vocab_size,
+    emb_dim=32,
+    latent_dim=32,
+    n_heads=4,
+    n_layers=2,
+    dim_feedforward=256,
+    max_pointers=population_size,
+)
+
+pointer = SelfAttentionPointer(
+    pointer_len=population_size,
+    d_model=32,
+)
+
+policy = DTSPolicy(
+    pop_to_vec_transformer=encoder,
+    pointer_transformer=pointer,
+    device="cpu",
+    train_every_n_gens=10,
+    learning_rate=2e-3,
+    final_lr=1e-3,
+    epsilon_greedy=1.0,
+    epsilon_greedy_decay=0.999,
+    min_epsilon=0.2,
+)
+
+dts = DeepTournamentSelection(policy, higher_is_better=True)
+```
+
+Use it as the EC-KitY selection method:
+
+```python
+selection_methods=[(dts, 1)]
+```
+
+Important parameters:
+
+- `population_size` determines the rank-embedding and pointer-table capacities.
+- `vocab_size` is the maximum integer gene value plus one.
+- `train_every_n_gens` controls how often accumulated trajectories train the policy.
+- `epsilon_greedy` and its decay control teacher-forced tournament selection versus learned selection.
+- `custom_reward_function`, when supplied, receives current fitness, previous fitness, and population arrays.
+- `device` can be `"cpu"` or `"cuda"` when a compatible PyTorch installation is available.
+
+## Fitness caching
+
+EC-KitY may evaluate unchanged vectors again across generations. `CachingEvaluator` avoids recomputing fitness for vectors it has already seen:
+
+```python
+from eckity_dts import CachingEvaluator
+
+evaluator = CachingEvaluator(MyEvaluator())
+print(evaluator.cache_stats())
+```
+
+## Compatibility
+
+- Python 3.9 or newer
+- EC-KitY 0.4.x
+- NumPy 2.0.2 or newer
+- SciPy 1.13.0 or newer
+- PyTorch 2.7.1 or newer
+- overrides 7.7.0 or newer
+
+These bounds are compatible with `eckity-dnc`, `eckity-bert-ga`, and `eckity-bert-gp`. None of the operator packages depends directly on another operator package.
+
+## Research repository
+
+The repository contains the full paper experiments for Graph Coloring, Set Cover, and TSP, together with benchmark instances, notebook-style runners, and result figures. These research resources are not included in the `eckity-dts` wheel.
+
+For repository development:
+
+```bash
+uv sync --extra dev --resolution lowest-direct
+uv run pytest
+```
+
+Experiment entry points remain available from a source checkout:
+
+```bash
+python -m deep_tournament_selection.experiments.graph_coloring --instance queen8_12.col.txt --generations 200
+python -m deep_tournament_selection.experiments.set_cover --instance scp41.txt --generations 200
+python -m deep_tournament_selection.experiments.tsp --instance att48.tsp --generations 200
+python run_experiments.py --selection both --runs 3 --generations 500
+```
+
+Results are written under `runs/`. Paper figures are stored under `figures/`, and the architecture diagram is under `images/`.
+
+## Development and release
+
+```bash
+uv run pytest
+uv run ruff check .
+uv build
+```
+
+Release preparation and manual PyPI upload commands are documented in [`RELEASING.md`](RELEASING.md).
+
+## License
+
+This project is licensed under the BSD 3-Clause License. See [`LICENSE`](LICENSE).

eckity_dts-0.1.0/deep_tournament_selection/__init__.py

@@ -0,0 +1,4 @@
+from .selection import DeepTournamentSelection, DTSPolicy
+from .caching_evaluator import CachingEvaluator
+
+__all__ = ["DeepTournamentSelection", "DTSPolicy", "CachingEvaluator"]

eckity_dts-0.1.0/deep_tournament_selection/caching_evaluator.py

@@ -0,0 +1,36 @@
+from collections import OrderedDict
+from eckity.evaluators.simple_individual_evaluator import SimpleIndividualEvaluator
+
+
+class CachingEvaluator(SimpleIndividualEvaluator):
+    def __init__(self, inner: SimpleIndividualEvaluator, max_size=None, events=None):
+        super().__init__(events=events)
+        self.inner = inner
+        self.max_size = max_size
+        self.fitness_cache = OrderedDict()
+        self.hits = 0
+        self.misses = 0
+
+    def evaluate_individual(self, individual):
+        key = tuple(individual.vector)
+        if key in self.fitness_cache:
+            self.hits += 1
+            self.fitness_cache.move_to_end(key)
+            return self.fitness_cache[key]
+
+        self.misses += 1
+        value = self.inner.evaluate_individual(individual)
+        self.fitness_cache[key] = value
+        if self.max_size is not None and len(self.fitness_cache) > self.max_size:
+            self.fitness_cache.popitem(last=False)
+        return value
+
+    def cache_stats(self):
+        total = self.hits + self.misses
+        hit_rate = self.hits / total if total else 0.0
+        return {
+            "hits": self.hits,
+            "misses": self.misses,
+            "size": len(self.fitness_cache),
+            "hit_rate": hit_rate,
+        }

eckity_dts-0.1.0/deep_tournament_selection/config.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass, field
+from typing import List
+
+
+@dataclass(frozen=True)
+class DTSConfig:
+    """Hyperparameters for the Deep Tournament Selection operator (paper setting)."""
+
+    tournament_size: int = 5
+    latent_dim: int = 32
+    dim_feedforward: int = 256
+    n_layers: int = 2
+    n_heads: int = 4
+    learning_rate: float = 2e-3
+    final_lr: float = 1e-3
+    train_every_n_gens: int = 10
+    epsilon_greedy: float = 1.0
+    epsilon_greedy_decay: float = 0.999
+    min_epsilon: float = 0.2
+
+
+POPULATION_SIZE = 100
+ELITISM = 2
+CROSSOVER_PROB = 0.5
+MUTATION_PROB = 0.5
+FLIP_MUTATION_PROB = 0.1
+RUNS = 15
+
+
+@dataclass(frozen=True)
+class TSPConfig:
+    instances: List[str] = field(
+        default_factory=lambda: [
+            "att48.tsp",
+            "berlin52.tsp",
+            "st70.tsp",
+            "lin105.tsp",
+            "pr107.tsp",
+            "pcb442.tsp",
+            "d1291.tsp",
+        ]
+    )
+    population_size: int = POPULATION_SIZE
+    generations: int = 1000
+    crossover_prob: float = CROSSOVER_PROB
+    mutation_prob: float = MUTATION_PROB
+    elitism: int = ELITISM
+    runs: int = RUNS
+
+
+@dataclass(frozen=True)
+class GraphColoringConfig:
+    instances: List[str] = field(
+        default_factory=lambda: [
+            "instances_games120.col.txt",
+            "instances_myciel7.col.txt",
+            "miles1000.col.txt",
+            "miles1500.col.txt",
+            "mulsol.i.2.col",
+            "queen8_12.col.txt",
+            "zeroin.i.1.col",
+            "zeroin.i.2.col",
+        ]
+    )
+    population_size: int = POPULATION_SIZE
+    generations: int = 6000
+    crossover_prob: float = CROSSOVER_PROB
+    mutation_prob: float = MUTATION_PROB
+    flip_mutation_prob: float = FLIP_MUTATION_PROB
+    penalty: float = 100.0
+    colors_margin: int = 10
+    elitism: int = ELITISM
+    runs: int = RUNS
+
+
+@dataclass(frozen=True)
+class SetCoverConfig:
+    instances: List[str] = field(
+        default_factory=lambda: [
+            "scp41.txt",
+            "scp51.txt",
+            "scp52.txt",
+            "scp53.txt",
+            "scp54.txt",
+            "scp56.txt",
+            "scp57.txt",
+            "scp64.txt",
+            "scp65.txt",
+        ]
+    )
+    population_size: int = POPULATION_SIZE
+    generations: int = 6000
+    crossover_prob: float = CROSSOVER_PROB
+    mutation_prob: float = MUTATION_PROB
+    flip_mutation_prob: float = FLIP_MUTATION_PROB
+    penalty: float = 100.0
+    weighted: bool = False
+    elitism: int = ELITISM
+    runs: int = RUNS

eckity_dts-0.1.0/deep_tournament_selection/elitist_breeder.py

@@ -0,0 +1,28 @@
+from eckity.breeders.simple_breeder import SimpleBreeder
+from eckity.genetic_operators.selections.elitism_selection import ElitismSelection
+
+
+class ElitistBreeder(SimpleBreeder):
+    def apply_breed(self, population):
+        for subpopulation in population.sub_populations:
+            nextgen_population = []
+
+            num_elites = subpopulation.n_elite
+            if num_elites > 0:
+                ElitismSelection(
+                    num_elites=num_elites,
+                    higher_is_better=subpopulation.higher_is_better,
+                ).apply_operator((subpopulation.individuals, nextgen_population))
+
+            selection = subpopulation.get_selection_methods()[0][0]
+            nextgen_population = selection.select(
+                subpopulation.individuals, nextgen_population
+            )
+            self.selected_individuals = nextgen_population
+
+            elites = nextgen_population[:num_elites]
+            offspring = self._apply_operators(
+                subpopulation.get_operators_sequence(),
+                nextgen_population[num_elites:],
+            )
+            subpopulation.individuals = elites + offspring

eckity_dts-0.1.0/deep_tournament_selection/logging_utils.py

@@ -0,0 +1,64 @@
+import json
+import os
+import time
+
+import numpy as np
+from eckity.statistics.statistics import Statistics
+
+
+class FileLogger(Statistics):
+    BASE_KEYS = ("mean", "std", "median", "max", "min", "time")
+
+    def __init__(
+        self,
+        output_path=None,
+        save_every_n_generations=50,
+        diversity_fn=None,
+        format_string=None,
+    ):
+        super().__init__(format_string or "")
+        self.output_path = output_path
+        self.save_every_n_generations = save_every_n_generations
+        self.diversity_fn = diversity_fn
+        keys = self.BASE_KEYS + (("population_diversity",) if diversity_fn else ())
+        self.generation_metrics = {key: [] for key in keys}
+        self._last_time = time.time()
+
+    def write_statistics(self, sender, data_dict):
+        sub_pop = data_dict["population"].sub_populations[0]
+        fitness_values = np.array(
+            [ind.get_pure_fitness() for ind in sub_pop.individuals], dtype=float
+        )
+        self.generation_metrics["mean"].append(float(np.mean(fitness_values)))
+        self.generation_metrics["std"].append(float(np.std(fitness_values)))
+        self.generation_metrics["median"].append(float(np.median(fitness_values)))
+        self.generation_metrics["max"].append(float(np.max(fitness_values)))
+        self.generation_metrics["min"].append(float(np.min(fitness_values)))
+
+        if self.diversity_fn is not None:
+            population = np.array([ind.vector for ind in sub_pop.individuals])
+            self.generation_metrics["population_diversity"].append(
+                float(self.diversity_fn(population))
+            )
+
+        now = time.time()
+        self.generation_metrics["time"].append(now - self._last_time)
+        self._last_time = now
+
+        if (
+            self.output_path
+            and self.save_every_n_generations
+            and data_dict["generation_num"] % self.save_every_n_generations == 0
+        ):
+            self.save()
+
+    def save(self, extra=None):
+        if self.output_path is None:
+            return None
+        os.makedirs(os.path.dirname(os.path.abspath(self.output_path)), exist_ok=True)
+        payload = dict(self.generation_metrics)
+        if extra:
+            payload.update(extra)
+        with open(self.output_path, "w") as f:
+            json.dump(payload, f, indent=2)
+        return self.output_path

eckity_dts-0.1.0/deep_tournament_selection/selection/__init__.py

@@ -0,0 +1,4 @@
+from .eckity_adapter import DeepTournamentSelection
+from .dts_policy import DTSPolicy
+
+__all__ = ["DeepTournamentSelection", "DTSPolicy"]

eckity_dts-0.1.0/deep_tournament_selection/selection/dts_aux.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+def get_ranks_from_fitness_values(fitness_values: np.ndarray):
+    order = np.argsort(-fitness_values, axis=1)
+    ranks = np.argsort(order, axis=1)
+    return ranks
+
+
+def get_ranking_based_rewards(fitness_values: np.ndarray, gamma=1.25):
+    fitness_values = fitness_values.reshape(1, -1)
+    ranks = get_ranks_from_fitness_values(fitness_values)
+    ranks_weights = 1 / (ranks + 1) ** gamma
+    rewards = ranks_weights * fitness_values
+    rewards = rewards.sum(axis=1)
+    return rewards
+
+
+def get_exponential_ranking_based_rewards(fitness_values: np.ndarray, alpha=0.2):
+    fitness_values = fitness_values.reshape(1, -1)
+    ranks = get_ranks_from_fitness_values(fitness_values)
+    ranks_weights = np.exp(-alpha * ranks)
+    rewards = ranks_weights * fitness_values
+    rewards = rewards.sum(axis=1)
+    return rewards
+
+
+def get_average_pairwise_hamming_distance(population: np.ndarray) -> float:
+    population = np.asarray(population)
+
+    if population.ndim != 2 or len(population) < 2:
+        return 0.0
+
+    diffs = population[:, None, :] != population[None, :, :]
+    pairwise_dist = diffs.sum(axis=2)
+    triu_indices = np.triu_indices(len(population), k=1)
+    return float(pairwise_dist[triu_indices].mean() / population.shape[1])
+
+
+def get_reward_from_fitness_scores(
+    cur_gen_fitness_values: np.ndarray,
+    prev_gen_fitness_values: np.ndarray,
+    population: np.ndarray | None = None,
+    top_k_to_consider: int = 5,
+    best_weight: float = 0.25,
+    top_k_weight: float = 1.0,
+    mean_weight: float = 0.25,
+    diversity_weight: float = 0.05,
+    eps: float = 1e-8,
+):
+    cur_gen_fitness_values = np.asarray(cur_gen_fitness_values, dtype=np.float64)
+    prev_gen_fitness_values = np.asarray(prev_gen_fitness_values, dtype=np.float64)
+
+    k = min(
+        top_k_to_consider, len(cur_gen_fitness_values), len(prev_gen_fitness_values)
+    )
+
+    def normalized_delta(cur_value: float, prev_value: float) -> float:
+        scale = max(abs(prev_value), eps)
+        return (cur_value - prev_value) / scale
+
+    cur_top_k_mean = np.partition(cur_gen_fitness_values, -k)[-k:].mean()
+    prev_top_k_mean = np.partition(prev_gen_fitness_values, -k)[-k:].mean()
+
+    cur_best = cur_gen_fitness_values.max()
+    prev_best = prev_gen_fitness_values.max()
+
+    cur_mean = cur_gen_fitness_values.mean()
+    prev_mean = prev_gen_fitness_values.mean()
+
+    diversity_reward = 0.0
+    if population is not None:
+        diversity_reward = get_average_pairwise_hamming_distance(population)
+
+    reward = (
+        top_k_weight * normalized_delta(cur_top_k_mean, prev_top_k_mean)
+        + best_weight * normalized_delta(cur_best, prev_best)
+        + mean_weight * normalized_delta(cur_mean, prev_mean)
+        + diversity_weight * diversity_reward
+    )
+    return float(reward)
+
+
+def get_trajectory_probability_from_log_probs(log_probs, selected_population_indices):
+    """
+    calculate total probability for the trajectory by summing the log probabilities over the selected indices
+    """
+    trajectory_prob = (
+        torch.gather(log_probs, 2, selected_population_indices.unsqueeze(-1))
+        .squeeze(-1)
+        .sum(dim=-1)
+    )
+    return trajectory_prob