Spaces:
Running
Running
| from .evolution import MolecularEvolution | |
| from .molecule import Molecule | |
| from core.predictors.mixture.mixture_dcn_predictor import MixtureDCNPredictor | |
| from core.base_fuel_library import BaseFuelLibrary | |
| from core.config import EvolutionConfig | |
| from typing import List, Tuple, Dict, Optional | |
| import numpy as np | |
| import wandb | |
| import pickle | |
| import torch | |
| from pathlib import Path | |
| import warnings | |
| from rdkit import RDLogger | |
| RDLogger.logger().setLevel(RDLogger.CRITICAL) | |
| warnings.filterwarnings("ignore") | |
| # Process-level singleton — loaded once, reused across all requests | |
| _dcn_predictor_instance: Optional[MixtureDCNPredictor] = None | |
| def get_dcn_predictor() -> MixtureDCNPredictor: | |
| global _dcn_predictor_instance | |
| if _dcn_predictor_instance is None: | |
| print("Initializing mixture DCN predictor...") | |
| _dcn_predictor_instance = MixtureDCNPredictor(max_models=5) | |
| _dcn_predictor_instance._initialize_models() | |
| return _dcn_predictor_instance | |
| class MixtureAwareMolecule(Molecule): | |
| """Extended Molecule class for mixture optimization with AD info.""" | |
| def __init__(self, *args, mixture_dcn=None, blend_ratio=None, | |
| ad_score=None, in_domain=None, mixture_ysi=None, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.mixture_dcn = mixture_dcn | |
| self.blend_ratio = blend_ratio | |
| self.ad_score = ad_score | |
| self.in_domain = in_domain | |
| self.mixture_ysi = mixture_ysi | |
| def to_dict(self): | |
| d = super().to_dict() | |
| d['mixture_dcn'] = self.mixture_dcn | |
| d['blend_ratio'] = self.blend_ratio | |
| d['ad_score'] = self.ad_score | |
| d['in_domain'] = self.in_domain | |
| d['mixture_ysi'] = self.mixture_ysi | |
| return d | |
| class MixturePredictionCache: | |
| """Cache DCN predictions for additive+base mixtures.""" | |
| def __init__(self, cache_file: str = "cache/mixture_dcn_cache.pkl"): | |
| self.cache_file = Path(cache_file) | |
| self.cache = self._load() | |
| def _load(self): | |
| if self.cache_file.exists(): | |
| try: | |
| with open(self.cache_file, 'rb') as f: | |
| cache = pickle.load(f) | |
| print(f" ✓ Loaded {len(cache)} cached DCN predictions") | |
| return cache | |
| except: | |
| return {} | |
| return {} | |
| def _save(self): | |
| self.cache_file.parent.mkdir(parents=True, exist_ok=True) | |
| with open(self.cache_file, 'wb') as f: | |
| pickle.dump(self.cache, f) | |
| def _key(self, smiles: str, base_fuel_type: str, fraction: float) -> str: | |
| return f"{smiles}|{base_fuel_type}|{fraction:.3f}" | |
| def get(self, additive_smiles: str, base_fuel_type: str, | |
| additive_fraction: float) -> float: | |
| return self.cache.get(self._key(additive_smiles, base_fuel_type, additive_fraction)) | |
| def get_batch(self, additive_smiles_list: list, base_fuel_type: str, | |
| additive_fraction: float) -> dict: | |
| return {smi: v for smi in additive_smiles_list | |
| if (v := self.get(smi, base_fuel_type, additive_fraction)) is not None} | |
| def set_batch(self, predictions: dict, base_fuel_type: str, | |
| additive_fraction: float): | |
| """Cache multiple predictions and flush to disk once.""" | |
| for smiles, dcn in predictions.items(): | |
| if dcn is not None: | |
| self.cache[self._key(smiles, base_fuel_type, additive_fraction)] = dcn | |
| self._save() | |
| def populate_from_database(self, csv_path: str, base_fuel_type: str, | |
| additive_fraction: float, base_smiles: List[str]): | |
| """Pre-load cache from mixture_database.csv using CN_Measured / CN_Regression.""" | |
| import pandas as pd | |
| from rdkit import Chem | |
| try: | |
| df = pd.read_csv(csv_path) | |
| except Exception: | |
| return | |
| def canonical(smi): | |
| try: | |
| return Chem.MolToSmiles(Chem.MolFromSmiles(smi)) if smi and str(smi) != 'nan' else None | |
| except Exception: | |
| return None | |
| base_canonical = {canonical(s) for s in base_smiles if s} | |
| frac_tol = 0.02 | |
| added = 0 | |
| for _, row in df.iterrows(): | |
| cn = row.get('CN_Measured') if not pd.isna(row.get('CN_Measured', float('nan'))) \ | |
| else row.get('CN_Regression') | |
| if pd.isna(cn): | |
| continue | |
| # Collect all (smiles, fraction) pairs from this row | |
| components = [] | |
| for k in range(1, 12): | |
| smi = canonical(str(row.get(f'fuel_{k}_smiles', '') or '')) | |
| frac = row.get(f'fraction_fuel_{k}') | |
| if smi and not pd.isna(frac): | |
| components.append((smi, float(frac))) | |
| if len(components) < 2: | |
| continue | |
| # Find if exactly one component is the additive (not a base component) | |
| for idx, (smi, frac) in enumerate(components): | |
| if smi in base_canonical: | |
| continue | |
| if abs(frac - additive_fraction) > frac_tol: | |
| continue | |
| rest = [c for j, c in enumerate(components) if j != idx] | |
| if all(c[0] in base_canonical for c in rest): | |
| key = self._key(smi, base_fuel_type, additive_fraction) | |
| if key not in self.cache: | |
| self.cache[key] = float(cn) | |
| added += 1 | |
| if added: | |
| self._save() | |
| print(f" ✓ Pre-loaded {added} DCN values from database") | |
| class MixtureAwareMolecularEvolution(MolecularEvolution): | |
| """ | |
| Mixture-aware evolution with Applicability Domain filtering. | |
| """ | |
| def __init__(self, config: EvolutionConfig, use_ad_filtering: bool = True): | |
| """ | |
| Args: | |
| config: Evolution configuration | |
| use_ad_filtering: Enable AD filtering (default: True) | |
| """ | |
| from .population import Population | |
| from core.predictors.pure_component.property_predictor import PropertyPredictor | |
| self.config = config | |
| self.predictor = PropertyPredictor(config) | |
| self.population = Population(config) | |
| self.uncertainty_filters = {} | |
| self.dcn_cache = MixturePredictionCache() | |
| # in-memory AD cache: smiles -> (ad_score, in_domain) | |
| self._ad_cache: Dict[str, Tuple[float, bool]] = {} | |
| self.mixture_predictor = get_dcn_predictor() | |
| self._load_base_fuel() | |
| self.use_ad_filtering = use_ad_filtering | |
| if use_ad_filtering: | |
| self._load_ad_checker() | |
| self._mol_db = None # lazy-loaded MolencoderDatabase cache | |
| # Pre-populate DCN cache from database | |
| mc = self.config.mixture_config | |
| db_path = Path(__file__).resolve().parent.parent.parent / "data" / "database" / "mixture_database.csv" | |
| self.dcn_cache.populate_from_database( | |
| str(db_path), | |
| mc.base_fuel_type, | |
| mc.additive_fraction, | |
| self.base_smiles, | |
| ) | |
| wandb.init(mode="disabled") | |
| def _load_base_fuel(self): | |
| """Load base fuel composition.""" | |
| mc = self.config.mixture_config | |
| if mc.base_fuel_smiles and mc.base_fuel_mole_fractions: | |
| self.base_smiles = mc.base_fuel_smiles | |
| self.base_fractions = mc.base_fuel_mole_fractions | |
| else: | |
| self.base_smiles, self.base_fractions = BaseFuelLibrary.get_base_fuel(mc.base_fuel_type) | |
| print(f"✓ Base fuel: {len(self.base_smiles)} components") | |
| def _load_ad_checker(self): | |
| """Load the trained One-Class SVM from HuggingFace Hub.""" | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| ad_path = hf_hub_download( | |
| repo_id="SalZa2004/mixture_ocvm_checker", | |
| filename="mixture_ocsvm.pkl", | |
| ) | |
| with open(ad_path, 'rb') as f: | |
| ad_data = pickle.load(f) | |
| self.svm = ad_data['svm'] | |
| self.scaler = ad_data['scaler'] | |
| print("✓ AD checker loaded") | |
| except Exception as e: | |
| print(f"⚠ AD checker could not be loaded ({e}) - disabling AD filtering") | |
| self.use_ad_filtering = False | |
| def _extract_mixture_embedding(self, additive_smiles: str) -> np.ndarray: | |
| from core.predictors.mixture.solvation_predictor.data.data import ( | |
| DataPoint, DatapointList, MolencoderDatabase, DataTensor | |
| ) | |
| mc = self.config.mixture_config | |
| base_ratio = 1.0 - mc.additive_fraction | |
| adjusted_base = [f * base_ratio for f in self.base_fractions] | |
| mixture_smiles = [additive_smiles] + self.base_smiles | |
| mixture_fractions = [mc.additive_fraction] + adjusted_base[:-1] # N-1 | |
| # Build DataPoint once — reused across all 10 models | |
| try: | |
| if self._mol_db is None: | |
| self._mol_db = MolencoderDatabase() | |
| mol_db = self._mol_db | |
| dp = DataPoint( | |
| smiles=mixture_smiles, | |
| targets=[0.0], | |
| features=[], | |
| molefracs=mixture_fractions, | |
| inp=self.mixture_predictor.args, | |
| mol_encoders=mol_db, | |
| ) | |
| data = DatapointList([dp]) | |
| args = self.mixture_predictor.args | |
| enc = dp.get_mol_encoder() | |
| if len(enc) < args.max_num_mols: | |
| enc += [enc[0]] * (args.max_num_mols - len(enc)) | |
| mol_encodings = [[enc[m]] for m in range(args.max_num_mols)] | |
| tensors = [DataTensor(mol_encodings[m], args, property=args.property) | |
| for m in range(args.max_num_mols)] | |
| except Exception: | |
| return None | |
| all_model_embeddings = [] | |
| for model in self.mixture_predictor.models: | |
| model.eval() | |
| captured = [] | |
| def hook_fn(module, input, output): | |
| captured.append(input[0].detach().cpu()) | |
| hook = model.ffn.register_forward_hook(hook_fn) | |
| try: | |
| with torch.no_grad(): | |
| _ = model(data, tensors) | |
| if captured: | |
| all_model_embeddings.append(captured[0].numpy().flatten()) | |
| except Exception: | |
| continue | |
| finally: | |
| hook.remove() | |
| if not all_model_embeddings: | |
| return None | |
| return np.mean(all_model_embeddings, axis=0) | |
| def _check_ad_batch(self, smiles_list: List[str]) -> Tuple[np.ndarray, np.ndarray]: | |
| """Check AD for a batch, using in-memory cache to skip already-seen SMILES.""" | |
| scores = np.full(len(smiles_list), -999.0) | |
| domain = np.zeros(len(smiles_list), dtype=bool) | |
| if not self.use_ad_filtering: | |
| return np.zeros(len(smiles_list)), np.ones(len(smiles_list), dtype=bool) | |
| # Split into cached vs needs-embedding | |
| need_emb_idx = [] | |
| for i, smi in enumerate(smiles_list): | |
| if smi in self._ad_cache: | |
| scores[i], domain[i] = self._ad_cache[smi] | |
| else: | |
| need_emb_idx.append(i) | |
| if not need_emb_idx: | |
| return scores, domain | |
| embeddings = [] | |
| valid_idx = [] | |
| for i in need_emb_idx: | |
| emb = self._extract_mixture_embedding(smiles_list[i]) | |
| if emb is not None: | |
| embeddings.append(emb) | |
| valid_idx.append(i) | |
| if embeddings: | |
| emb_arr = self.scaler.transform(np.array(embeddings)) | |
| scores_v = self.svm.decision_function(emb_arr) | |
| domain_v = self.svm.predict(emb_arr) == 1 | |
| for j, idx in enumerate(valid_idx): | |
| scores[idx] = scores_v[j] | |
| domain[idx] = domain_v[j] | |
| self._ad_cache[smiles_list[idx]] = (float(scores_v[j]), bool(domain_v[j])) | |
| return scores, domain | |
| def predict_mixture_ysi(self, additive_smiles: str) -> Optional[float]: | |
| """ | |
| Predict mixture YSI using the linear blending law (mass-fraction weighted). | |
| Args: | |
| additive_smiles: SMILES of the additive molecule. | |
| Returns: | |
| ysi_mix: Predicted mixture YSI, or None if any step fails. | |
| """ | |
| from core.blending.blending_law import blend_ysi_mass_weighted | |
| mc = self.config.mixture_config | |
| base_ratio = 1.0 - mc.additive_fraction | |
| all_smiles = [additive_smiles] + self.base_smiles | |
| all_mole_fracs = [mc.additive_fraction] + [f * base_ratio for f in self.base_fractions] | |
| if 'ysi' not in self.predictor.predictors: | |
| from core.predictors.pure_component.generic import GenericPredictor | |
| from core.predictors.pure_component.hf_models import load_models | |
| paths = load_models() | |
| self.predictor.predictors['ysi'] = GenericPredictor(paths['ysi'], 'YSI') | |
| props = self.predictor.predict_all_properties(all_smiles) | |
| ysi_values = props.get('ysi', []) | |
| if len(ysi_values) != len(all_smiles): | |
| return None | |
| return blend_ysi_mass_weighted(all_smiles, all_mole_fracs, ysi_values) | |
| def predict_mixture_ysi_batch(self, smiles_list: List[str]) -> List[Optional[float]]: | |
| """Predict mixture YSI for a list of additives.""" | |
| return [self.predict_mixture_ysi(smi) for smi in smiles_list] | |
| def _create_molecules(self, smiles_list: List[str]) -> Tuple[List[MixtureAwareMolecule], Dict]: | |
| if not smiles_list: | |
| return [], { | |
| 'total': 0, | |
| 'cn_none': 0, | |
| 'ysi_none': 0, | |
| 'ad_filtered': 0, | |
| 'passed': 0, | |
| } | |
| mc = self.config.mixture_config | |
| filter_stats = { | |
| 'total': len(smiles_list), | |
| 'cn_none': 0, | |
| 'ysi_none': 0, | |
| 'ad_filtered': 0, | |
| 'passed': 0, | |
| } | |
| # STEP 1: Check AD for all molecules (returns numpy arrays) | |
| ad_scores_array, in_domain_array = self._check_ad_batch(smiles_list) | |
| n_in_domain = in_domain_array.sum() | |
| n_out_domain = (~in_domain_array).sum() | |
| print(f" → AD Check: {n_in_domain} in-domain, {n_out_domain} out-of-domain") | |
| print(f" Score range: [{ad_scores_array.min():.3f}, {ad_scores_array.max():.3f}]") | |
| # Short-circuit: nothing passes AD, skip expensive predictions | |
| if n_in_domain == 0: | |
| filter_stats['ad_filtered'] = len(smiles_list) | |
| return [], filter_stats | |
| # STEP 2: Get DCN predictions only for in-domain molecules (cache-first) | |
| in_domain_indices = [i for i, ok in enumerate(in_domain_array) if ok] | |
| in_domain_smiles = [smiles_list[i] for i in in_domain_indices] | |
| cached_dcns = self.dcn_cache.get_batch( | |
| in_domain_smiles, mc.base_fuel_type, mc.additive_fraction | |
| ) | |
| uncached_smiles = [s for s in in_domain_smiles if s not in cached_dcns] | |
| if uncached_smiles: | |
| new_dcns = self.mixture_predictor.predict_batch_mixtures( | |
| additive_smiles_list=uncached_smiles, | |
| base_smiles=self.base_smiles, | |
| base_mole_fractions=self.base_fractions, | |
| additive_fraction=mc.additive_fraction, | |
| verbose=False, | |
| ) | |
| self.dcn_cache.set_batch( | |
| dict(zip(uncached_smiles, new_dcns)), | |
| mc.base_fuel_type, | |
| mc.additive_fraction, | |
| ) | |
| cached_dcns.update(zip(uncached_smiles, new_dcns)) | |
| mixture_dcns_in_domain = {orig_idx: cached_dcns.get(smiles_list[orig_idx]) | |
| for orig_idx in in_domain_indices} | |
| # STEP 3: Predict YSI only for in-domain additives + base components | |
| # Ensure YSI predictor is loaded | |
| if 'ysi' not in self.predictor.predictors: | |
| from core.predictors.pure_component.generic import GenericPredictor | |
| from core.predictors.pure_component.hf_models import load_models | |
| paths = load_models() | |
| self.predictor.predictors['ysi'] = GenericPredictor(paths['ysi'], 'YSI') | |
| base_ratio = 1.0 - mc.additive_fraction | |
| all_unique_smiles = in_domain_smiles + self.base_smiles | |
| props_all = self.predictor.predict_all_properties(all_unique_smiles) | |
| ysi_for_in_domain = props_all.get('ysi', [None] * len(all_unique_smiles)) | |
| # Build a lookup by original index for in-domain additives | |
| ysi_by_orig_idx = {orig_idx: ysi_for_in_domain[pos] | |
| for pos, orig_idx in enumerate(in_domain_indices)} | |
| # Pre-compute base component YSI (same for every additive) | |
| base_ysi = ysi_for_in_domain[len(in_domain_smiles):] | |
| from core.blending.blending_law import blend_ysi_mass_weighted | |
| # STEP 4: Create molecules — only iterate over in-domain candidates | |
| filter_stats['ad_filtered'] = n_out_domain | |
| molecules = [] | |
| for i in in_domain_indices: | |
| smiles = smiles_list[i] | |
| dcn = mixture_dcns_in_domain.get(i) | |
| if dcn is None: | |
| filter_stats['cn_none'] += 1 | |
| continue | |
| additive_ysi = ysi_by_orig_idx.get(i) | |
| mole_fracs = [mc.additive_fraction] + [f * base_ratio for f in self.base_fractions] | |
| mixture_ysi = blend_ysi_mass_weighted( | |
| [smiles] + self.base_smiles, | |
| mole_fracs, | |
| [additive_ysi] + list(base_ysi), | |
| ) | |
| filter_stats['passed'] += 1 | |
| molecules.append(MixtureAwareMolecule( | |
| smiles=smiles, | |
| cn=dcn, | |
| cn_error=abs(dcn - mc.target_mixture_dcn), | |
| cn_score=dcn, | |
| mixture_dcn=dcn, | |
| blend_ratio=mc.additive_fraction, | |
| ad_score=float(ad_scores_array[i]), | |
| in_domain=bool(in_domain_array[i]), | |
| mixture_ysi=mixture_ysi, | |
| ysi=mixture_ysi # mirrors mixture_ysi so Population NSGA-II can use it | |
| )) | |
| return molecules, filter_stats | |
| def _log_generation_stats(self, generation: int): | |
| """Log stats with AD metrics.""" | |
| mols = self.population.molecules | |
| n_invalid = sum(1 for m in mols if not m.chemical_valid) | |
| if self.config.maximize_cn: | |
| best = max(mols, key=lambda m: m.cn) | |
| avg_metric = np.mean([m.cn for m in mols]) | |
| best_metric = best.cn | |
| else: | |
| best = min(mols, key=lambda m: m.cn_error) | |
| avg_metric = np.mean([m.cn_error for m in mols]) | |
| best_metric = best.cn_error | |
| avg_ratio = np.mean([m.blend_ratio for m in mols if m.blend_ratio is not None]) | |
| # NEW: AD stats | |
| if self.use_ad_filtering and mols: | |
| avg_ad = np.mean([m.ad_score for m in mols if m.ad_score is not None]) | |
| n_in_domain = sum(1 for m in mols if m.in_domain) | |
| else: | |
| avg_ad = 0.0 | |
| n_in_domain = len(mols) | |
| # Console | |
| pareto_size = len(self.population.pareto_front()) if self.config.minimize_ysi else 0 | |
| extra = f" | Pareto: {pareto_size}" if self.config.minimize_ysi else "" | |
| if self.use_ad_filtering: | |
| print( | |
| f"Gen {generation}/{self.config.generations} | " | |
| f"Pop {len(mols)} | " | |
| f"Best: {best_metric:.3f} | " | |
| f"AD: {avg_ad:.3f}{extra}" | |
| ) | |
| else: | |
| print( | |
| f"Gen {generation}/{self.config.generations} | " | |
| f"Pop {len(mols)} | " | |
| f"Best: {best_metric:.3f} | " | |
| f"Invalid: {n_invalid}{extra}" | |
| ) | |
| # W&B scalars | |
| log_dict = { | |
| "generation": generation, | |
| "population_size": len(mols), | |
| "best_mixture_dcn" if self.config.maximize_cn else "best_dcn_error": best_metric, | |
| "avg_mixture_dcn" if self.config.maximize_cn else "avg_dcn_error": avg_metric, | |
| "invalid_fraction": n_invalid / len(mols) if mols else 0, | |
| "avg_blend_ratio": avg_ratio, | |
| } | |
| if self.use_ad_filtering: | |
| log_dict["avg_ad_score"] = avg_ad | |
| log_dict["in_domain_fraction"] = n_in_domain / len(mols) if mols else 0 | |
| if self.config.minimize_ysi: | |
| ysi_vals = [m.mixture_ysi for m in mols if m.mixture_ysi is not None] | |
| log_dict["avg_mixture_ysi"] = np.mean(ysi_vals) if ysi_vals else float('nan') | |
| log_dict["pareto_front_size"] = pareto_size | |
| wandb.log(log_dict) | |
| # W&B table every 10 gens | |
| if generation % 6 == 0: | |
| if self.use_ad_filtering: | |
| table_data = [ | |
| [m.smiles, m.mixture_dcn, m.cn_error, m.mixture_ysi, m.ad_score, m.in_domain] | |
| for m in sorted(mols, key=lambda x: x.cn_error)[:50] | |
| ] | |
| columns = ["SMILES", "DCN", "Error", "YSI", "AD Score", "In Domain"] | |
| else: | |
| table_data = [ | |
| [m.smiles, m.mixture_dcn, m.cn_error, m.mixture_ysi, m.blend_ratio] | |
| for m in sorted(mols, key=lambda x: x.cn_error)[:50] | |
| ] | |
| columns = ["SMILES", "DCN", "Error", "YSI", "Blend Ratio"] | |
| table = wandb.Table(data=table_data, columns=columns) | |
| wandb.log({f"top_molecules_gen_{generation}": table}) | |
| def initialize_population(self, initial_smiles: List[str]) -> int: | |
| """Initialize population with AD-based filtering.""" | |
| print("Predicting properties for initial population...") | |
| molecules, filter_stats = self._create_molecules(initial_smiles) | |
| if filter_stats['total'] > 0: | |
| print(f"\n Initial filtering: {filter_stats['total']} → {filter_stats['passed']} passed") | |
| print(f" AD filtered: {filter_stats['ad_filtered']} | " | |
| f"CN None: {filter_stats['cn_none']}") | |
| return self.population.add_molecules(molecules) | |
| def _generate_offspring(self, survivors: List[MixtureAwareMolecule]) -> Tuple[List[MixtureAwareMolecule], Dict]: | |
| """Generate offspring using AD-based filtering (no Tanimoto / uncertainty filters).""" | |
| import random | |
| from rdkit import Chem | |
| target_count = self.config.population_size - len(survivors) | |
| max_attempts = target_count * self.config.max_offspring_attempts | |
| all_children: List[str] = [] | |
| new_molecules: List[MixtureAwareMolecule] = [] | |
| cumulative_stats = { | |
| 'total': 0, | |
| 'cn_none': 0, | |
| 'ysi_none': 0, | |
| 'ad_filtered': 0, | |
| 'passed': 0, | |
| } | |
| print(f" → Generating offspring (target: {target_count})...") | |
| for _ in range(max_attempts): | |
| if len(new_molecules) >= target_count: | |
| break | |
| weights = np.array([m.fitness(self.config) for m in survivors]) | |
| weights /= weights.sum() | |
| parent = survivors[np.random.choice(len(survivors), p=weights)] | |
| mol = Chem.MolFromSmiles(parent.smiles) | |
| if mol is None: | |
| continue | |
| children = self._mutate_molecule(mol) | |
| all_children.extend(children[:self.config.mutations_per_parent]) | |
| if len(all_children) >= self.config.batch_size: | |
| batch_mols, batch_stats = self._create_molecules(all_children) | |
| new_molecules.extend(batch_mols) | |
| for key in cumulative_stats: | |
| cumulative_stats[key] += batch_stats.get(key, 0) | |
| all_children = [] | |
| if all_children: | |
| batch_mols, batch_stats = self._create_molecules(all_children) | |
| new_molecules.extend(batch_mols) | |
| for key in cumulative_stats: | |
| cumulative_stats[key] += batch_stats.get(key, 0) | |
| print(f" ✓ Generated {len(new_molecules)} valid offspring") | |
| print(f" Filtering: {cumulative_stats['total']} → {cumulative_stats['passed']} | " | |
| f"AD filtered: {cumulative_stats['ad_filtered']} | " | |
| f"CN None: {cumulative_stats['cn_none']} " | |
| f"(property constraints applied at end)") | |
| return new_molecules, cumulative_stats | |
| def _predict_densities(self, smiles_list: List[str]) -> Dict[str, Optional[float]]: | |
| """Return a SMILES→density map, handling featurization drop-outs.""" | |
| from core.shared_features import featurize_df | |
| if not smiles_list: | |
| return {} | |
| # Start with all None; only successful predictions overwrite entries. | |
| density_map: Dict[str, Optional[float]] = {smi: None for smi in smiles_list} | |
| # return_df=True gives back which SMILES actually survived featurization, | |
| # avoiding a silent length mismatch when descriptors fail for some molecules. | |
| result = featurize_df(smiles_list, return_df=True) | |
| if result is None or result[0] is None: | |
| return density_map | |
| X, df_valid = result | |
| valid_smiles = df_valid['SMILES'].tolist() | |
| preds = self.predictor.predictors['density'].predict_from_features(X) | |
| for smi, val in zip(valid_smiles, preds): | |
| try: | |
| density_map[smi] = float(val) if val is not None and np.isfinite(val) else None | |
| except (TypeError, ValueError): | |
| density_map[smi] = None | |
| n_none = sum(1 for v in density_map.values() if v is None) | |
| if n_none: | |
| smi_sample = next(s for s in smiles_list if density_map.get(s) is None) | |
| print(f" ⚠ Density None for {n_none}/{len(smiles_list)} SMILES " | |
| f"(e.g. '{smi_sample}')") | |
| return density_map | |
| def _compute_mixture_bp(self, additive_smiles_list: List[str]) -> Dict[str, Optional[float]]: | |
| """Compute Riazi-Daubert mixture boiling point (°C) for each unique additive SMILES.""" | |
| from core.blending.blending_law import blend_bp_riazi_daubert | |
| mc = self.config.mixture_config | |
| base_ratio = 1.0 - mc.additive_fraction | |
| unique = list(dict.fromkeys(additive_smiles_list)) | |
| density_map = self._predict_densities(unique + self.base_smiles) | |
| base_densities = [density_map.get(smi) for smi in self.base_smiles] | |
| mole_fracs = [mc.additive_fraction] + [f * base_ratio for f in self.base_fractions] | |
| results = {} | |
| for smi in unique: | |
| all_densities = [density_map.get(smi)] + base_densities | |
| # blend_bp_riazi_daubert expects specific gravity in g/cm³; | |
| # the density predictor returns kg/m³, so divide by 1000. | |
| all_densities_gcc = [ | |
| d / 1000.0 if d is not None else None for d in all_densities | |
| ] | |
| results[smi] = blend_bp_riazi_daubert( | |
| [smi] + self.base_smiles, | |
| mole_fracs, | |
| all_densities_gcc, | |
| ) | |
| return results | |
| def _compute_mixture_density(self, additive_smiles_list: List[str]) -> Dict[str, Optional[float]]: | |
| """Compute mixture density (g/cm³) for each unique additive SMILES using Eq. 3 (β_ij = 0).""" | |
| from core.blending.blending_law import blend_density | |
| mc = self.config.mixture_config | |
| base_ratio = 1.0 - mc.additive_fraction | |
| unique = list(dict.fromkeys(additive_smiles_list)) | |
| density_map = self._predict_densities(unique + self.base_smiles) | |
| base_densities = [density_map.get(smi) for smi in self.base_smiles] | |
| mole_fracs = [mc.additive_fraction] + [f * base_ratio for f in self.base_fractions] | |
| results = {} | |
| for smi in unique: | |
| all_densities = [density_map.get(smi)] + base_densities | |
| # Density predictor returns kg/m³; convert to g/cm³ for blending law | |
| all_densities_gcc = [ | |
| d / 1000.0 if d is not None else None for d in all_densities | |
| ] | |
| result_gcc = blend_density( | |
| [smi] + self.base_smiles, | |
| mole_fracs, | |
| all_densities_gcc, | |
| ) | |
| results[smi] = result_gcc * 1000.0 if result_gcc is not None else None | |
| return results | |
| def _sort_df(self, df): | |
| """Sort without a hard cn_error cutoff. | |
| The base class applies cn_error < 5 or < 15 thresholds that are fine for | |
| pure-component evolution but too tight for mixture DCN predictions, which | |
| can sit further from the target early in the run. | |
| """ | |
| if df.empty: | |
| return df | |
| if self.config.maximize_cn: | |
| if self.config.minimize_ysi and "ysi" in df.columns: | |
| return df.sort_values(["cn", "ysi"], ascending=[False, True]) | |
| return df.sort_values("cn", ascending=False) | |
| else: | |
| if self.config.minimize_ysi and "ysi" in df.columns: | |
| return df.sort_values(["cn_error", "ysi"], ascending=True) | |
| return df.sort_values("cn_error", ascending=True) | |
| def _apply_mixture_filters(self, df): | |
| """Filter final_df by mixture_bp and mixture_density bounds. | |
| Molecules where a property is None (couldn't be computed) are left in — | |
| only molecules with a computed value that falls outside the range are removed. | |
| """ | |
| import pandas as pd | |
| if df.empty: | |
| return df | |
| mask = pd.Series(True, index=df.index) | |
| for prop, (lo, hi) in self.config.mixture_filters.items(): | |
| if prop not in df.columns: | |
| continue | |
| col = df[prop] | |
| if lo is not None: | |
| mask &= col.isna() | (col >= lo) | |
| if hi is not None: | |
| mask &= col.isna() | (col <= hi) | |
| filtered = df[mask].copy() | |
| removed = len(df) - len(filtered) | |
| if removed > 0: | |
| print(f" Mixture property filters removed {removed}/{len(df)} molecules") | |
| filtered["rank"] = range(1, len(filtered) + 1) | |
| return filtered | |
| def _generate_results(self): | |
| """Generate final DataFrames and append mixture_bp and mixture_density.""" | |
| final_df, pareto_df, unfiltered_df = super()._generate_results() | |
| # Collect smiles from all three DataFrames so bp/density can be computed | |
| # even when final_df is still empty before mixture filtering. | |
| all_smiles = [] | |
| for df in (final_df, pareto_df, unfiltered_df): | |
| if not df.empty and 'smiles' in df.columns: | |
| all_smiles.extend(df['smiles'].tolist()) | |
| if all_smiles: | |
| unique_smiles = list(dict.fromkeys(all_smiles)) | |
| print(" Computing mixture boiling points (Riazi-Daubert)...") | |
| bp_map = self._compute_mixture_bp(unique_smiles) | |
| print(" Computing mixture densities (Eq. 3, β_ij = 0)...") | |
| density_map = self._compute_mixture_density(unique_smiles) | |
| for result_df in (final_df, pareto_df, unfiltered_df): | |
| if not result_df.empty and 'smiles' in result_df.columns: | |
| result_df['mixture_bp'] = result_df['smiles'].map(bp_map) | |
| result_df['mixture_density'] = result_df['smiles'].map(density_map) | |
| # Apply mixture-specific property constraints now that blended columns exist. | |
| # This is done here rather than in the base class because mixture_bp and | |
| # mixture_density don't exist until after the blending computation above. | |
| final_df = self._apply_mixture_filters(final_df) | |
| return final_df, pareto_df, unfiltered_df |