Hybrid native backend, protein-aware diffusion, and smarter candidate triage - #14
Closed
cosmic-hydra with Codex wants to merge 2 commits into
Closed
Hybrid native backend, protein-aware diffusion, and smarter candidate triage#14cosmic-hydra with Codex wants to merge 2 commits into
cosmic-hydra with Codex wants to merge 2 commits into
Conversation
Co-authored-by: cosmic-hydra <140935487+cosmic-hydra@users.noreply.github.com>
Codex created this pull request from a session on behalf of
cosmic-hydra
April 19, 2026 10:17
View session
cosmic-hydra
marked this pull request as ready for review
April 19, 2026 10:17
Review Summary by QodoHybrid native backend, protein-aware diffusion, and intelligent candidate triage
WalkthroughsDescription• Native C++/CUDA physics backend with Lennard-Jones energy, forces, and FEP computation • Protein-aware diffusion model with cross-attention and FiLM conditioning for ligand generation • Joint protein–ligand interaction model producing binding affinity and contact maps • Smarter candidate selection combining MC-dropout uncertainty, EHVI, diversity, and retrosynthesis feasibility • Time-based dataset splitting and MC-dropout uncertainty quantification for predictions • Training stability improvements: target normalization, calibration tracking, and energy regularization Diagramflowchart LR
A["Native Physics Backend<br/>C++/CUDA"] -->|compute_energy<br/>compute_forces<br/>run_fep| B["MD Simulator<br/>& Energy Calc"]
C["Protein Context<br/>Embeddings"] -->|cross-attention<br/>FiLM| D["Diffusion Model<br/>Ligand Generation"]
E["Ligand & Residue<br/>Features"] -->|joint encoder| F["Interaction Head<br/>Affinity + Contacts"]
G["Candidates<br/>with Uncertainty"] -->|EHVI + Diversity<br/>+ Retrosynthesis| H["Ranked Selection<br/>Top-K + Explorers"]
B -->|coordinates| F
D -->|generated molecules| G
F -->|binding scores| H
File Changes1. drug_discovery/native/__init__.py
|
Code Review by Qodo
1. calculate_energy return shape
|
![qodo-code-review[bot]](https://avatars.githubusercontent.com/in/484649?s=60&v=4){kind=small}
Comment on lines
352
to
355
| except Exception as e: | ||
| logger.error(f"Energy calculation error: {e}") | ||
| return None | ||
|
|
There was a problem hiding this comment.
1. Calculate_energy return shape 🐞 Bug ≡ Correctness
EnergyCalculator.calculate_energy() returns plain None in the exception path even when return_coords=True, so callers that unpack (energy, coords) will crash.
Agent Prompt
### Issue description
`EnergyCalculator.calculate_energy(..., return_coords=True)` is unpacked by callers, but the exception handler returns a bare `None`, causing an unpacking crash.
### Issue Context
The method already returns `(None, None)` for several early failure cases when `return_coords=True`; the remaining exception handler should follow the same contract.
### Fix Focus Areas
- drug_discovery/physics/md_simulator.py[306-355]
### Suggested fix
- In the `except Exception as e:` block, return `(None, None)` when `return_coords` is `True`, else return `None`.
- Keep logging behavior unchanged.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
408
to
409
| except Exception as e: | ||
| logger.error(f"Geometry optimization error: {e}") |
There was a problem hiding this comment.
2. Optimize_geometry return shape 🐞 Bug ≡ Correctness
EnergyCalculator.optimize_geometry() returns a 2-tuple (None, None) on exception even when return_coords=True, so callers expecting 3 values will crash.
Agent Prompt
### Issue description
`EnergyCalculator.optimize_geometry(..., return_coords=True)` is unpacked as 3 values, but the exception path returns only 2 values.
### Issue Context
The function already returns `(None, None, None)` for some early failures when `return_coords=True`; the exception handler should do the same.
### Fix Focus Areas
- drug_discovery/physics/md_simulator.py[356-410]
### Suggested fix
- In the `except Exception as e:` block, return `(None, None, None)` when `return_coords` is `True`, else return `(None, None)`.
- Ensure type annotation remains accurate.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+275
to
+282
| def _energy_penalty(batch): | ||
| coords = getattr(batch, "pos", None) | ||
| if coords is None: | ||
| return None | ||
| try: | ||
| return torch.mean(compute_energy(coords.to(self.device), reduce=False)) | ||
| except Exception: | ||
| return None |
There was a problem hiding this comment.
3. Energy penalty mixes molecules 🐞 Bug ≡ Correctness
Pipeline energy regularization computes Lennard-Jones energy over concatenated PyG batch.pos, introducing cross-molecule interactions and an incorrect penalty signal during graph training.
Agent Prompt
### Issue description
Energy regularization currently treats a PyG minibatch as a single molecule/system, so the LJ penalty is mathematically wrong and mixes unrelated graphs.
### Issue Context
`batch.pos` in PyG is concatenated across graphs; `batch.batch` (or `batch.ptr`) provides the graph membership needed to split coordinates.
### Fix Focus Areas
- drug_discovery/pipeline.py[274-283]
- drug_discovery/native/backend.py[55-71]
### Suggested fix
- In `_energy_penalty`, if `batch` has `batch`/`ptr`, split `coords` into per-graph tensors and compute energy per graph, then average.
- Example approach: iterate graphs using `batch.ptr` slices; pad to `(B, Nmax, 3)` and call `compute_energy(..., reduce=False)`; or call `compute_energy` per graph in a loop if batches are small.
- If graph membership is unavailable, return `None` (don’t apply penalty).
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+43
to
+54
| def predict_with_uncertainty(self, features: torch.Tensor, samples: int = 8) -> tuple[np.ndarray, np.ndarray]: | ||
| """Predict with MC dropout-based uncertainty.""" | ||
| preds = [] | ||
| was_training = self.model.training | ||
| self.model.train() | ||
| with torch.no_grad(): | ||
| for _ in range(max(1, samples)): | ||
| preds.append(self.model(features.to(self.device)).cpu()) | ||
| self.model.eval() | ||
| if was_training: | ||
| self.model.train() | ||
| stacked = torch.stack(preds, dim=0) |
There was a problem hiding this comment.
4. Mc dropout mutates batchnorm 🐞 Bug ≡ Correctness
predict_with_uncertainty() sets the whole model to train() for MC dropout, which updates BatchNorm running stats and can silently change subsequent predictions and training behavior.
Agent Prompt
### Issue description
MC-dropout inference should not update BatchNorm running statistics; current implementation does.
### Issue Context
The GNN uses `nn.BatchNorm1d`, which updates buffers in `train()` mode.
### Fix Focus Areas
- drug_discovery/evaluation/predictor.py[43-57]
### Suggested fix
- Keep the model in `eval()` mode, then enable dropout layers only:
- Save current mode.
- Call `self.model.eval()`.
- Recursively set `Dropout` modules to `train()` (or use a helper that toggles only dropout).
- Run multiple forward passes.
- Restore original module modes.
- Alternatively, if you must call `train()`, set all BatchNorm modules back to `eval()` before sampling.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+241
to
+260
| if cond_ctx is not None: | ||
| with torch.no_grad(): | ||
| if torch.rand(1, device=self.device).item() < self._uncond_prob: | ||
| eff_cond_ctx = None | ||
| eff_cond_mask = None | ||
| else: | ||
| eff_cond_ctx = cond_ctx | ||
| eff_cond_mask = cond_mask | ||
| uncond_pos, uncond_atom = self.model( | ||
| atom_types, | ||
| pos, | ||
| edge_index, | ||
| t, | ||
| batch, | ||
| protein_context=eff_cond_ctx, | ||
| protein_mask=eff_cond_mask, | ||
| ) | ||
| eps_pos = uncond_pos + g_scale * (cond_pos - uncond_pos) | ||
| eps_atom = uncond_atom + g_scale * (cond_atom - uncond_atom) | ||
| else: |
There was a problem hiding this comment.
5. Cfg guidance randomly disabled 🐞 Bug ≡ Correctness
Diffusion sampling computes the “unconditional” pass by sometimes reusing the conditioned context, which frequently makes uncond==cond and reduces classifier-free guidance to a no-op.
Agent Prompt
### Issue description
Classifier-free guidance requires both a conditioned and an unconditional prediction at each step; current sampling sometimes computes the 'unconditional' prediction with the conditioning still applied.
### Issue Context
`uncond_dropout_prob` is typically used during training to randomly drop conditioning. During sampling, CFG usually computes:
- `eps_cond = model(x, cond)`
- `eps_uncond = model(x, None)`
- `eps = eps_uncond + s*(eps_cond-eps_uncond)`
### Fix Focus Areas
- drug_discovery/models/diffusion_generator.py[226-262]
### Suggested fix
- In `sample()`, when `cond_ctx` is not None, always compute `uncond_*` with `protein_context=None` (and mask None/zeros) rather than probabilistically.
- Keep `uncond_dropout_prob` for training-time conditioning dropout (if/when training code is added).
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+207
to
+211
| if ctx.dim() == 2: | ||
| ctx = ctx.unsqueeze(0).expand(num_molecules, -1, -1) | ||
| elif ctx.dim() == 3 and ctx.size(0) != num_molecules: | ||
| ctx = ctx.expand(num_molecules, -1, -1) | ||
| protein_mask = torch.zeros(ctx.shape[:2], dtype=torch.bool, device=self.device) |
There was a problem hiding this comment.
6. Protein_context expand crash 🐞 Bug ☼ Reliability
_prepare_condition() uses Tensor.expand to match num_molecules, which throws when protein_context has a batch dimension that is neither 1 nor num_molecules.
Agent Prompt
### Issue description
`Tensor.expand` only works for singleton dimensions; mismatched batch sizes will crash sampling.
### Issue Context
Users may call `sample(num_molecules=N, protein_context=ctx)` with ctx already batched.
### Fix Focus Areas
- drug_discovery/models/diffusion_generator.py[203-212]
### Suggested fix
- If `ctx.dim()==3`:
- If `ctx.size(0)==1`, `expand` is fine.
- If `ctx.size(0)==num_molecules`, keep as-is.
- Otherwise: raise a clear ValueError explaining the mismatch, or use `repeat` only if semantics are intended.
- Add a unit test or assertion to cover mismatch behavior.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
There was a problem hiding this comment.
Pull request overview
This PR upgrades ZANE toward a hybrid system by adding a native (Torch-extension) physics backend, extending diffusion generation with protein conditioning + classifier-free guidance, and introducing smarter candidate triage plus new evaluation/splitting utilities.
Changes:
- Add a lazily-built Torch C++ extension backend (
compute_energy,compute_forces,run_fep) and integrate it into MD simulation and training hooks. - Extend diffusion sampling/modeling with protein/pocket conditioning and add a protein–ligand interaction model.
- Add candidate selection utilities (uncertainty + EHVI + diversity + retrosynthesis filtering), plus time-based splits and an enrichment factor metric.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 17 comments.
Show a summary per file
| File | Description |
|---|---|
| setup.py | Packages native C++ source into the distribution. |
| drug_discovery/training/trainer.py | Adds target normalization, calibration logging, and an energy regularization hook. |
| drug_discovery/pipeline.py | Adds time split support, energy-regularized training wiring, and new candidate triage logic. |
| drug_discovery/physics/md_simulator.py | Integrates native energy/forces/FEP and propagates coordinate tensors. |
| drug_discovery/physics/init.py | Re-exports native physics functions via the physics API. |
| drug_discovery/optimization/selection.py | New candidate selector combining uncertainty, EHVI, and diversity. |
| drug_discovery/optimization/init.py | Exposes new selection utilities from the optimization package. |
| drug_discovery/native/forcefield.cpp | Implements the Torch extension entrypoints (energy/forces/FEP). |
| drug_discovery/native/backend.py | Python wrapper that lazy-loads/compiles the extension with Torch fallbacks. |
| drug_discovery/native/init.py | Public native module exports. |
| drug_discovery/models/protein_ligand.py | New protein–ligand interaction model (affinity + contact map). |
| drug_discovery/models/diffusion_generator.py | Adds protein conditioning and classifier-free guidance to diffusion sampling. |
| drug_discovery/models/init.py | Exposes the new protein–ligand model/config. |
| drug_discovery/evaluation/predictor.py | Adds MC-dropout uncertainty APIs and enrichment factor metric. |
| drug_discovery/data/dataset.py | Adds chronological (time-based) dataset split helper. |
| README.md | Documents new native backend, protein-aware diffusion, interaction head, and triage. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+89
to
+95
| self.config.ehvi_weight * base_metric + self.config.uncertainty_weight * unc + self.config.diversity_weight * div | ||
| ) | ||
| metric_arr = np.asarray(metric_values, dtype=np.float32) | ||
| ehvi = expected_hypervolume_improvement(metric_arr, reference_point=self.config.reference_point) | ||
| ranked = [] | ||
| for idx, score in enumerate(scores): | ||
| combined = float(score + 0.1 * ehvi[idx]) |
There was a problem hiding this comment.
The linear score uses base_metric directly (ehvi_weight * base_metric), but EHVI is later computed from both base_metric and qed_score. This double-counts base_metric while qed_score only affects the much smaller 0.1 * ehvi term; consider making the weighting consistent with the intended multi-objective ranking.
Suggested change
| self.config.ehvi_weight * base_metric + self.config.uncertainty_weight * unc + self.config.diversity_weight * div | |
| ) | |
| metric_arr = np.asarray(metric_values, dtype=np.float32) | |
| ehvi = expected_hypervolume_improvement(metric_arr, reference_point=self.config.reference_point) | |
| ranked = [] | |
| for idx, score in enumerate(scores): | |
| combined = float(score + 0.1 * ehvi[idx]) | |
| self.config.uncertainty_weight * unc + self.config.diversity_weight * div | |
| ) | |
| metric_arr = np.asarray(metric_values, dtype=np.float32) | |
| ehvi = expected_hypervolume_improvement(metric_arr, reference_point=self.config.reference_point) | |
| ranked = [] | |
| for idx, score in enumerate(scores): | |
| combined = float(score + self.config.ehvi_weight * ehvi[idx]) |
Comment on lines
+209
to
+210
| elif ctx.dim() == 3 and ctx.size(0) != num_molecules: | ||
| ctx = ctx.expand(num_molecules, -1, -1) |
There was a problem hiding this comment.
ctx.expand(num_molecules, -1, -1) will fail unless ctx.size(0) is 1. If protein_context is already batched (B!=1) and num_molecules differs, this will throw at runtime. Validate that ctx.size(0) is 1 or num_molecules, or tile using repeat when appropriate.
Suggested change
| elif ctx.dim() == 3 and ctx.size(0) != num_molecules: | |
| ctx = ctx.expand(num_molecules, -1, -1) | |
| elif ctx.dim() == 3: | |
| if ctx.size(0) == num_molecules: | |
| pass | |
| elif ctx.size(0) == 1: | |
| ctx = ctx.expand(num_molecules, -1, -1) | |
| else: | |
| raise ValueError( | |
| f"protein_context batch size must be 1 or match num_molecules " | |
| f"({num_molecules}), got {ctx.size(0)}" | |
| ) | |
| else: | |
| raise ValueError( | |
| f"protein_context must have 2 or 3 dimensions, got {ctx.dim()}" | |
| ) |
Comment on lines
+242
to
+256
| with torch.no_grad(): | ||
| if torch.rand(1, device=self.device).item() < self._uncond_prob: | ||
| eff_cond_ctx = None | ||
| eff_cond_mask = None | ||
| else: | ||
| eff_cond_ctx = cond_ctx | ||
| eff_cond_mask = cond_mask | ||
| uncond_pos, uncond_atom = self.model( | ||
| atom_types, | ||
| pos, | ||
| edge_index, | ||
| t, | ||
| batch, | ||
| protein_context=eff_cond_ctx, | ||
| protein_mask=eff_cond_mask, |
There was a problem hiding this comment.
Classifier-free guidance here is stochastic: the "unconditional" pass sometimes uses the full condition (eff_cond_ctx = cond_ctx), making guidance ineffective and adding unrelated randomness. Guidance should deterministically combine a true unconditional pass (no condition) with the conditional pass; if you want to save compute, make that an explicit option instead of random dropout during sampling.
Suggested change
| with torch.no_grad(): | |
| if torch.rand(1, device=self.device).item() < self._uncond_prob: | |
| eff_cond_ctx = None | |
| eff_cond_mask = None | |
| else: | |
| eff_cond_ctx = cond_ctx | |
| eff_cond_mask = cond_mask | |
| uncond_pos, uncond_atom = self.model( | |
| atom_types, | |
| pos, | |
| edge_index, | |
| t, | |
| batch, | |
| protein_context=eff_cond_ctx, | |
| protein_mask=eff_cond_mask, | |
| uncond_pos, uncond_atom = self.model( | |
| atom_types, | |
| pos, | |
| edge_index, | |
| t, | |
| batch, | |
| protein_context=None, | |
| protein_mask=None, |
Comment on lines
+97
to
+102
| cursor = torch.zeros((batch_size,), device=features.device, dtype=torch.long) | ||
| for idx, b in enumerate(batch): | ||
| pos = cursor[b].item() | ||
| dense[b, pos] = features[idx] | ||
| padding_mask[b, pos] = False | ||
| cursor[b] += 1 |
There was a problem hiding this comment.
_dense_from_batch loops over batch and calls .item() each iteration to index cursor. If batch is on GPU this causes per-element CPU/GPU synchronization and will be a major bottleneck. Prefer a vectorized approach (e.g., torch_geometric.utils.to_dense_batch or scatter-based indexing).
Suggested change
| cursor = torch.zeros((batch_size,), device=features.device, dtype=torch.long) | |
| for idx, b in enumerate(batch): | |
| pos = cursor[b].item() | |
| dense[b, pos] = features[idx] | |
| padding_mask[b, pos] = False | |
| cursor[b] += 1 | |
| if batch.numel() > 0: | |
| perm = torch.argsort(batch, stable=True) | |
| sorted_batch = batch[perm] | |
| sorted_idx = torch.arange(batch.numel(), device=batch.device, dtype=torch.long) | |
| group_start = torch.ones_like(sorted_batch, dtype=torch.bool) | |
| group_start[1:] = sorted_batch[1:] != sorted_batch[:-1] | |
| start_idx = torch.where(group_start, sorted_idx, torch.zeros_like(sorted_idx)) | |
| start_idx = torch.cummax(start_idx, dim=0).values | |
| pos_sorted = sorted_idx - start_idx | |
| pos = torch.empty_like(pos_sorted) | |
| pos[perm] = pos_sorted | |
| dense[batch, pos] = features | |
| padding_mask[batch, pos] = False |
Comment on lines
+43
to
+53
| def predict_with_uncertainty(self, features: torch.Tensor, samples: int = 8) -> tuple[np.ndarray, np.ndarray]: | ||
| """Predict with MC dropout-based uncertainty.""" | ||
| preds = [] | ||
| was_training = self.model.training | ||
| self.model.train() | ||
| with torch.no_grad(): | ||
| for _ in range(max(1, samples)): | ||
| preds.append(self.model(features.to(self.device)).cpu()) | ||
| self.model.eval() | ||
| if was_training: | ||
| self.model.train() |
There was a problem hiding this comment.
Switching the entire model to train() for MC dropout can update BatchNorm (and other training-mode state) even under no_grad, changing the model during inference. Prefer enabling dropout layers only while keeping BatchNorm in eval (or otherwise freezing running stats).
Suggested change
| def predict_with_uncertainty(self, features: torch.Tensor, samples: int = 8) -> tuple[np.ndarray, np.ndarray]: | |
| """Predict with MC dropout-based uncertainty.""" | |
| preds = [] | |
| was_training = self.model.training | |
| self.model.train() | |
| with torch.no_grad(): | |
| for _ in range(max(1, samples)): | |
| preds.append(self.model(features.to(self.device)).cpu()) | |
| self.model.eval() | |
| if was_training: | |
| self.model.train() | |
| def _enable_mc_dropout(self) -> list[tuple[torch.nn.Module, bool]]: | |
| """Enable training mode only for dropout layers and return their previous states.""" | |
| dropout_types = ( | |
| torch.nn.Dropout, | |
| torch.nn.Dropout1d, | |
| torch.nn.Dropout2d, | |
| torch.nn.Dropout3d, | |
| torch.nn.AlphaDropout, | |
| torch.nn.FeatureAlphaDropout, | |
| ) | |
| dropout_states = [] | |
| for module in self.model.modules(): | |
| if isinstance(module, dropout_types): | |
| dropout_states.append((module, module.training)) | |
| module.train() | |
| return dropout_states | |
| def predict_with_uncertainty(self, features: torch.Tensor, samples: int = 8) -> tuple[np.ndarray, np.ndarray]: | |
| """Predict with MC dropout-based uncertainty.""" | |
| preds = [] | |
| features = features.to(self.device) | |
| was_training = self.model.training | |
| self.model.eval() | |
| dropout_states = self._enable_mc_dropout() | |
| with torch.no_grad(): | |
| for _ in range(max(1, samples)): | |
| preds.append(self.model(features).cpu()) | |
| for module, module_was_training in dropout_states: | |
| module.train(module_was_training) | |
| self.model.train(was_training) |
Comment on lines
+172
to
+174
| if backend is not None: | ||
| try: | ||
| delta_f = backend.run_fep(ligand_coords, protein_coords, lambda_schedule, sigma, epsilon) |
There was a problem hiding this comment.
lambda_schedule can be None here, but the C++ extension binding expects a torch::Tensor. Passing None will raise and force the torch fallback every time. Convert None to an empty tensor (so the extension can apply its default schedule), or update the binding to accept an optional tensor.
Suggested change
| if backend is not None: | |
| try: | |
| delta_f = backend.run_fep(ligand_coords, protein_coords, lambda_schedule, sigma, epsilon) | |
| native_lambda_schedule = ( | |
| lambda_schedule | |
| if lambda_schedule is not None | |
| else torch.empty(0, device=ligand_coords.device, dtype=ligand_coords.dtype) | |
| ) | |
| if backend is not None: | |
| try: | |
| delta_f = backend.run_fep( | |
| ligand_coords, protein_coords, native_lambda_schedule, sigma, epsilon | |
| ) |
Comment on lines
+110
to
+114
| batch = [] | ||
| for b_idx in range(dense.size(0)): | ||
| n = int(keep[b_idx].sum().item()) | ||
| batch.extend([b_idx] * n) | ||
| batch_tensor = torch.tensor(batch, device=dense.device, dtype=torch.long) |
There was a problem hiding this comment.
_flatten_from_dense builds batch_tensor via a Python list and then torch.tensor(...). This is slow for large batches and can incur host->device copies. Consider constructing batch_tensor with repeat_interleave using per-batch counts to keep it vectorized and device-friendly.
Suggested change
| batch = [] | |
| for b_idx in range(dense.size(0)): | |
| n = int(keep[b_idx].sum().item()) | |
| batch.extend([b_idx] * n) | |
| batch_tensor = torch.tensor(batch, device=dense.device, dtype=torch.long) | |
| counts = keep.sum(dim=1, dtype=torch.long) | |
| batch_tensor = torch.repeat_interleave( | |
| torch.arange(dense.size(0), device=dense.device, dtype=torch.long), | |
| counts, | |
| ) |
Comment on lines
+43
to
+45
| def predict_with_uncertainty(self, features: torch.Tensor, samples: int = 8) -> tuple[np.ndarray, np.ndarray]: | ||
| """Predict with MC dropout-based uncertainty.""" | ||
| preds = [] |
There was a problem hiding this comment.
New MC-dropout APIs (predict_with_uncertainty / predict_from_smiles_with_uncertainty) are introduced here, but there are no tests exercising them (e.g., output shapes, restoring model mode, stochasticity). Since tests/test_predictor_comprehensive.py already covers other predictor methods, please add coverage for these new uncertainty methods as well.
Comment on lines
+304
to
+312
| """Compute enrichment factor at given top fraction.""" | ||
| if len(y_true) == 0: | ||
| return 0.0 | ||
| n_top = max(1, int(len(y_true) * top_fraction)) | ||
| order = np.argsort(-y_score.reshape(-1)) | ||
| top_idx = order[:n_top] | ||
| hits_top = float(np.sum(y_true.reshape(-1)[top_idx] > 0)) | ||
| hit_rate = hits_top / n_top | ||
| baseline = float(np.mean(y_true.reshape(-1) > 0)) if np.mean(y_true) != 0 else 1e-8 |
There was a problem hiding this comment.
enrichment_factor() introduces new evaluation logic but appears untested. Add unit tests for key edge cases (all-zero labels, all-one labels, ties, very small arrays) to ensure the metric is stable and matches the intended definition.
Suggested change
| """Compute enrichment factor at given top fraction.""" | |
| if len(y_true) == 0: | |
| return 0.0 | |
| n_top = max(1, int(len(y_true) * top_fraction)) | |
| order = np.argsort(-y_score.reshape(-1)) | |
| top_idx = order[:n_top] | |
| hits_top = float(np.sum(y_true.reshape(-1)[top_idx] > 0)) | |
| hit_rate = hits_top / n_top | |
| baseline = float(np.mean(y_true.reshape(-1) > 0)) if np.mean(y_true) != 0 else 1e-8 | |
| """Compute enrichment factor at given top fraction. | |
| The enrichment factor is defined as the hit rate among the highest-scoring | |
| compounds divided by the overall hit rate in the full set. | |
| Edge cases are handled explicitly to keep the metric stable: | |
| - empty inputs return 0.0 | |
| - datasets with no positive labels return 0.0 | |
| - ties are ordered deterministically via a stable sort | |
| - very small arrays still evaluate at least one top-ranked item | |
| """ | |
| y_true_flat = np.asarray(y_true).reshape(-1) | |
| y_score_flat = np.asarray(y_score).reshape(-1) | |
| if y_true_flat.size == 0 or y_score_flat.size == 0: | |
| return 0.0 | |
| if y_true_flat.size != y_score_flat.size: | |
| raise ValueError("y_true and y_score must have the same number of elements") | |
| n_samples = y_true_flat.size | |
| n_top = min(n_samples, max(1, int(n_samples * top_fraction))) | |
| positives = y_true_flat > 0 | |
| baseline = float(np.mean(positives)) | |
| if baseline == 0.0: | |
| return 0.0 | |
| order = np.argsort(-y_score_flat, kind="mergesort") | |
| top_idx = order[:n_top] | |
| hits_top = float(np.sum(positives[top_idx])) | |
| hit_rate = hits_top / n_top |
| if self._target_mean is None or self._target_std is None: | ||
| self._target_mean = targets.mean() | ||
| std = targets.std() | ||
| self._target_std = std if std > 1e-6 else torch.tensor(1.0, device=targets.device) |
There was a problem hiding this comment.
std is a 0-d Tensor here, so std > 1e-6 produces a Tensor boolean and the ternary will raise "Boolean value of Tensor is ambiguous" at runtime. Use std.item() for the comparison or clamp std (e.g., std = std.clamp_min(1e-6)).
Suggested change
| self._target_std = std if std > 1e-6 else torch.tensor(1.0, device=targets.device) | |
| self._target_std = std.clamp_min(1e-6) |
Owner
|
@codex[agent] @claude[agent] please resolve conflicts and commit changes |
cosmic-hydra
marked this pull request as draft
April 19, 2026 10:27
Co-authored-by: cosmic-hydra <140935487+cosmic-hydra@users.noreply.github.com>
Copilot stopped work on behalf of
cosmic-hydra due to an error
April 19, 2026 10:31
Copilot stopped work on behalf of
cosmic-hydra due to an error
April 19, 2026 10:34
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Upgrades ZANE into the requested hybrid system: native physics backend, GPU-aware diffusion with protein conditioning, joint interaction modeling, and smarter candidate selection with retrosynthesis-aware filtering.
Native physics backend
compute_energy,compute_forces,run_fep) with CPU fallback and packaging support.Protein-aware generative modeling
Candidate triage and evaluation
Training stability and regularization
Example usage (physics-backed energy in training):