Spaces:
Sleeping
Sleeping
| from typing import List | |
| from core.config import EvolutionConfig | |
| from .molecule import Molecule | |
| import pandas as pd | |
| import numpy as np | |
| from pymoo.util.nds.non_dominated_sorting import NonDominatedSorting | |
| def _crowding_distance(F: np.ndarray) -> np.ndarray: | |
| """Compute NSGA-II crowding distance for an (n, m) objective matrix.""" | |
| n, m = F.shape | |
| dist = np.zeros(n) | |
| for obj in range(m): | |
| idx = np.argsort(F[:, obj]) | |
| dist[idx[0]] = dist[idx[-1]] = np.inf | |
| spread = F[idx[-1], obj] - F[idx[0], obj] | |
| if spread == 0: | |
| continue | |
| for k in range(1, n - 1): | |
| dist[idx[k]] += (F[idx[k + 1], obj] - F[idx[k - 1], obj]) / spread | |
| return dist | |
| class Population: | |
| """Manages the population of molecules.""" | |
| def __init__(self, config: EvolutionConfig): | |
| self.config = config | |
| self.molecules: List[Molecule] = [] | |
| self.seen_smiles: set = set() | |
| def add_molecule(self, mol: Molecule) -> bool: | |
| """Add a molecule if not already in the population.""" | |
| if mol.smiles in self.seen_smiles: | |
| return False | |
| self.molecules.append(mol) | |
| self.seen_smiles.add(mol.smiles) | |
| return True | |
| def add_molecules(self, molecules: List[Molecule]) -> int: | |
| """Add multiple molecules, return count added.""" | |
| return sum(self.add_molecule(mol) for mol in molecules) | |
| def _objective_matrix(self, molecules: List[Molecule]) -> np.ndarray: | |
| """Build (n, 2) minimisation objective matrix. | |
| pymoo always minimises, so CN is negated when maximising. | |
| """ | |
| cn_col = ( | |
| np.array([-m.cn for m in molecules]) | |
| if self.config.maximize_cn | |
| else np.array([m.cn_error for m in molecules]) | |
| ) | |
| ysi_col = np.array([m.ysi for m in molecules]) | |
| return np.column_stack([cn_col, ysi_col]) | |
| def pareto_front(self) -> List[Molecule]: | |
| """Return Pareto-front molecules using pymoo non-dominated sorting.""" | |
| if not self.config.minimize_ysi or not self.molecules: | |
| return [] | |
| F = self._objective_matrix(self.molecules) | |
| front0 = NonDominatedSorting().do(F, only_non_dominated_front=True) | |
| return [self.molecules[i] for i in front0] | |
| def get_survivors(self) -> List[Molecule]: | |
| """Select survivors via NSGA-II rank + crowding-distance truncation.""" | |
| target_size = int(self.config.population_size * self.config.survivor_fraction) | |
| if not self.config.minimize_ysi: | |
| return sorted( | |
| self.molecules, | |
| key=lambda m: self.config.cn_objective(m.cn), | |
| reverse=True, | |
| )[:target_size] | |
| F = self._objective_matrix(self.molecules) | |
| fronts = NonDominatedSorting().do(F) | |
| survivor_indices: List[int] = [] | |
| for front in fronts: | |
| slots_left = target_size - len(survivor_indices) | |
| if len(front) <= slots_left: | |
| survivor_indices.extend(front.tolist()) | |
| else: | |
| # Partial front: keep most diverse by crowding distance | |
| crowding = _crowding_distance(F[front]) | |
| best = front[np.argsort(-crowding)[:slots_left]] | |
| survivor_indices.extend(best.tolist()) | |
| break | |
| return [self.molecules[i] for i in survivor_indices] | |
| def to_dataframe(self) -> pd.DataFrame: | |
| """Convert population to DataFrame.""" | |
| df = pd.DataFrame([m.to_dict() for m in self.molecules]) | |
| if self.config.maximize_cn: | |
| sort_cols = ["cn", "ysi"] if self.config.minimize_ysi else ["cn"] | |
| ascending = [False, True] if self.config.minimize_ysi else False | |
| else: | |
| sort_cols = ["cn_error", "ysi"] if self.config.minimize_ysi else ["cn_error"] | |
| ascending = True | |
| df = df.sort_values(sort_cols, ascending=ascending) | |
| df.insert(0, "rank", range(1, len(df) + 1)) | |
| return df | |