diff --git a/CHANGELOG.md b/CHANGELOG.md index 15390237..65acf54e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ Changelog * Raise an error when the `annihilate` or `decouple` lambda schedule is used with a perturbable molecule created by `sire.morph.annihilate()`, which also removes the bonded terms. ABFE ligands should be created with `sire.morph.decouple()`, leaving the choice between decoupling and annihilation to the lambda schedule [#239](https://github.com/OpenBioSim/somd2/pull/239). * Only fall back to the Aldeghi Boresch restraint search protocol when the system can't be used with the RXRX protocol. Systems that can't be used with either protocol, or restraint search settings that would save too few frames, now raise an error before the restraint search trajectory is run, rather than after a misleading fallback warning [#241](https://github.com/OpenBioSim/somd2/pull/241). * Save the GCMC end-state topologies, which include the additional waters, at the first checkpoint of each run as well as the last, and keep them on restart, so that they match the trajectory and checkpoints while a simulation is running, if it crashes, or if a restart never reaches a checkpoint [#244](https://github.com/OpenBioSim/somd2/pull/244). +* Add `gcmc_batch_size` and `gcmc_num_attempts` options to set the number of GCMC trials evaluated in parallel per batch and in total per move, which were previously fixed at the defaults of the GCMC sampler. Both can be changed on restart, e.g. to reduce wasted trials when more than one is accepted per move [#246](https://github.com/OpenBioSim/somd2/pull/246). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/README.md b/README.md index 43925081..3f0009a7 100644 --- a/README.md +++ b/README.md @@ -419,6 +419,13 @@ related options, run: somd2 --help | grep -A2 ' --gcmc' ``` +Each GCMC move evaluates its trials in parallel batches, set with +`--gcmc-batch-size`, up to a total of `--gcmc-num-attempts`. Only the first +accepted trial of a batch is used, so if more than one trial is accepted per +move on average, the rest of each batch is wasted, and the batch size should be +reduced. The [viewer](#viewer) shows the number accepted per move for each λ +window. These options can be changed on restart. + > [!IMPORTANT] > GCMC is only supported when using the CUDA or OpenCL platforms. @@ -754,6 +761,10 @@ The viewer shows: end state from the latest saved coordinates of its window, with the mutated residues and any ligand in detail. Peptides with fewer than five amino acids are depicted like any other perturbed molecule. +- In either of the above, the waters within 6 Å of the perturbed molecule or + mutated residues, which can be turned off. For GCMC simulations, the waters + in the GCMC region are also shown, along with the region itself, leaving out + ghost waters. - The MBAR free energy, PMF, overlap matrix, and forward and backward convergence. These are updated in the background as new data is written. For ABFE simulations, the standard state correction for the Boresch restraint is @@ -763,7 +774,8 @@ The viewer shows: - Energy components as a function of time for each λ window. These are written at each checkpoint, or at every energy sample when using `--save-energy-components`. -- GCMC and terminal flip Monte Carlo statistics, if active. +- GCMC and terminal flip Monte Carlo statistics, if active, including the + number of waters in the GCMC region over time. - Any restraints, e.g. a Boresch restraint for an ABFE simulation, or Morse restraints for a ring-breaking perturbation. - The λ schedule, REST2 scale factors, configuration options, and tuned PME diff --git a/src/somd2/_viewer/_binding_site.py b/src/somd2/_viewer/_binding_site.py index 766e1ae5..30339d89 100644 --- a/src/somd2/_viewer/_binding_site.py +++ b/src/somd2/_viewer/_binding_site.py @@ -31,6 +31,8 @@ import numpy as _np +from ._coordinates import Mismatch as _Mismatch + # Protein residues with an atom within this distance, in nm, of a perturbed # molecule are shown in detail. _CUTOFF = 0.6 @@ -44,13 +46,18 @@ # are within this distance, in nm, of a mutated residue. _NEARBY = 1.2 -# The topologies read most recently, which are kept to a limited number. +# Waters are given their own chain. +_WATER_CHAIN = "W" + +# The topologies read most recently, and the layouts of their waters, which +# are kept to a limited number. _topologies = _OrderedDict() +_water_layouts = _OrderedDict() _topologies_lock = _threading.Lock() _MAX_TOPOLOGIES = 16 -def binding_site(topology0, topology1, saved, saved1=None): +def binding_site(topology0, topology1, saved, saved1=None, gcmc=None): """ The protein and perturbed molecules, or a mutated protein, at the saved coordinates. @@ -68,6 +75,11 @@ def binding_site(topology0, topology1, saved, saved1=None): The coordinates from the λ = 1 window, for that end state of a mutated protein. Those from the λ = 0 window are used if None. + gcmc: dict + The GCMC region, as the "selection" of atoms at its centre and its + "radius" in nm, which are None if it is the whole box, or None if + GCMC wasn't used. + Returns ------- @@ -79,19 +91,26 @@ def binding_site(topology0, topology1, saved, saved1=None): hydrogens, and any other molecules, e.g. a ligand, at each end state, with the serial numbers of the atoms of the mutated residues, of the residues near them or the other molecules, and of the other molecules. - None if there is nothing to show. + Also the waters near the perturbed molecules or mutated residues, with + their serial numbers, and with GCMC, those in its region, leaving out + ghosts, and the region's centre and radius in Å. None if there is + nothing to show. """ topology = _topology(topology0, topology1) if topology is None: return None + layout = _water_layout( + topology0, topology1, gcmc["selection"] if gcmc is not None else None + ) if topology["kind"] == "mutation": - return _mutation(topology, saved, saved1 or saved) + return _mutation(topology, saved, saved1 or saved, layout, gcmc) proteins = [_coordinates(saved, topology, m) for m in topology["proteins"]] ligands = [_coordinates(saved, topology, m) for m in topology["others"]] _image(saved.box, proteins, ligands) near = _np.concatenate(ligands) + region = _waters(layout, gcmc, topology, saved, near) lines = [] pocket = [] serial = 0 @@ -115,6 +134,7 @@ def binding_site(topology0, topology1, saved, saved1=None): serial += 1 ligand_atoms.append(serial) lines.append(_ligand_atom(serial, record, i, xyz)) + waters, serial = _water_records(region, serial, lines) lines.append("END") return { @@ -122,6 +142,9 @@ def binding_site(topology0, topology1, saved, saved1=None): "pdb": "\n".join(lines) + "\n", "pocket": pocket, "ligand": ligand_atoms, + "waters": waters, + "num_waters": 0 if region is None else len(region[1]), + "gcmc": _region_summary(region, gcmc), "time_ps": saved.time_ps, } @@ -135,11 +158,11 @@ def site_kind(topology0, topology1): return topology["kind"] if topology is not None else None -def _mutation(topology, saved0, saved1): +def _mutation(topology, saved0, saved1, layout=None, gcmc=None): """ - The mutated protein and any other molecules at each end state, each from - its own window, with the λ = 1 window superposed on the λ = 0 window by - the protein's Cα atoms. + The mutated protein, any other molecules and nearby waters at each end + state, each from its own window, with the λ = 1 window superposed on the + λ = 0 window by the protein's Cα atoms. """ def ca(saved): @@ -163,10 +186,6 @@ def ca(saved): _image(saved.box, proteins, others) frames.append((proteins, others)) - if saved1 is not saved0: - fit = _superpose(ca(saved1), ca(saved0)) - frames[1] = tuple([fit(x) for x in group] for group in frames[1]) - def mutated_positions(state): proteins = frames[state][0] return _np.concatenate( @@ -204,6 +223,29 @@ def mutated_positions(state): for residues, molecule in zip(nearby, topology["proteins"]): residues -= molecule["mutated"] + # Waters near the mutated residues, or the other molecules shown. + regions = [ + _waters( + layout, + gcmc, + topology, + saved, + _np.concatenate([mutated_positions(state)] + [others[i] for i in shown]), + ) + for state, (saved, (_, others)) in enumerate(zip((saved0, saved1), frames)) + ] + + if saved1 is not saved0: + fit = _superpose(ca(saved1), ca(saved0)) + frames[1] = tuple([fit(x) for x in group] for group in frames[1]) + if regions[1] is not None: + centre, waters, atoms = regions[1] + regions[1] = ( + None if centre is None else fit(centre[None, :])[0], + fit(waters.reshape(-1, 3)).reshape(waters.shape), + atoms, + ) + states = [] for state, (proteins, others) in enumerate(frames): lines = [] @@ -231,6 +273,7 @@ def mutated_positions(state): serial += 1 ligand.append(serial) lines.append(_ligand_atom(serial, record, i, xyz)) + waters, serial = _water_records(regions[state], serial, lines) lines.append("END") states.append( { @@ -238,6 +281,9 @@ def mutated_positions(state): "mutated": mutated, "nearby": neighbours, "ligand": ligand, + "waters": waters, + "num_waters": 0 if regions[state] is None else len(regions[state][1]), + "gcmc": _region_summary(regions[state], gcmc), } ) @@ -278,6 +324,106 @@ def _image(box, proteins, others): positions -= shift +def _waters(layout, gcmc, topology, saved, near): + """ + The centre of any GCMC region, and the atoms, by water, of the waters near + the given atoms, in their image, and with GCMC, those in its region, + leaving out ghosts, with the names and elements of each water's atoms. + None if there are no waters, or with GCMC, the ghosts aren't known. + """ + if layout is None or (gcmc is not None and saved.ghost_waters is None): + return None + starts = saved.molecule_starts(layout["waters"], topology["num_system_atoms"]) + atom_indices = starts[:, None] + _np.array([a[0] for a in layout["atoms"]]) + if atom_indices.size and atom_indices.max() >= len(saved.positions): + raise _Mismatch() + indices = starts + layout["oxygen"] + unwrapped = saved.positions[indices] + oxygens = unwrapped + + reference = near.mean(axis=0) + periodic = _is_periodic(saved.box) + if periodic: + box = _np.asarray(saved.box, dtype=float) + inverse = _np.linalg.inv(box) + + centre = None + if gcmc is not None and layout["selection"]: + selected = _np.concatenate( + [ + _coordinates(saved, topology, molecule, atoms=atoms) + for molecule, atoms in layout["selection"] + ] + ) + # As in GCMC, each atom is taken in the image nearest to the first. + if periodic: + offsets = selected - selected[0] + selected = selected[0] + offsets - _np.round(offsets @ inverse) @ box + centre = selected.mean(axis=0) + + if periodic: + offsets = oxygens - reference + oxygens = reference + offsets - _np.round(offsets @ inverse) @ box + if centre is not None: + centre = centre - _np.round((centre - reference) @ inverse) @ box + + present = _np.ones(len(oxygens), dtype=bool) + if saved.ghost_waters is not None: + # Ghosts beyond the waters in the topology are of one written before + # the additional GCMC waters were added to it. + ghosts = _np.asarray(saved.ghost_waters, dtype=int) + present[ghosts[ghosts < len(present)]] = False + + # Only the waters that could be close enough are checked against each atom. + spread = _np.linalg.norm(near - reference, axis=1).max() + candidates = _np.where( + present & (_np.linalg.norm(oxygens - reference, axis=1) < spread + _CUTOFF) + )[0] + distances = _np.linalg.norm( + oxygens[candidates][:, None, :] - near[None, :, :], axis=-1 + ).min(axis=1, initial=_np.inf) + shown = _np.zeros(len(oxygens), dtype=bool) + shown[candidates[distances < _CUTOFF]] = True + if centre is not None: + shown |= present & (_np.linalg.norm(oxygens - centre, axis=1) < gcmc["radius"]) + + # Each water's atoms are moved into the image of its oxygen. + shifts = (oxygens - unwrapped)[shown] + atoms = saved.positions[atom_indices[shown]] + shifts[:, None, :] + return centre, atoms, layout["atoms"] + + +def _water_records(region, serial, lines): + """ + Add HETATM records for the atoms of the waters shown to lines, returning + their serial numbers and the last serial number. + """ + serials = [] + if region is not None: + _, waters, atoms = region + for i, positions in enumerate(waters): + for (_, name, element), xyz in zip(atoms, positions): + serial += 1 + serials.append(serial) + record = (name, "HOH", i + 1, element) + lines.append(_atom("HETATM", serial, record, _WATER_CHAIN, xyz)) + return serials, serial + + +def _region_summary(region, gcmc): + """ + The centre and radius of a GCMC region in Å, with the selection of atoms + at its centre, or None. + """ + if region is None or region[0] is None: + return None + return { + "centre": (10.0 * region[0]).tolist(), + "radius": 10.0 * gcmc["radius"], + "selection": gcmc["selection"], + } + + def _superpose(mobile, target): """ The function that best fits a set of coordinates onto another, by @@ -424,6 +570,98 @@ def entry(index, protein): return topology +def _water_layout(topology0, topology1, selection=None): + """ + The waters, as the index of each one's first atom, in the order that GCMC + numbers them, the position, name and element of each one's oxygen and + hydrogens, and, given a GCMC selection, the atoms at the centre of its + region, by molecule. Read once for each pair of topologies and selection. + None if there are no waters. + """ + from ._summary import _file_key + + key = (_file_key(topology0), _file_key(topology1), selection) + with _topologies_lock: + if key in _water_layouts: + _water_layouts.move_to_end(key) + return _water_layouts[key] + + import sire as _sr + + from ._data import sire_lock + + def waters(system): + try: + numbers = {mol.number() for mol in system.molecules("water")} + except KeyError: + return set() + return { + i for i, mol in enumerate(system.molecules()) if mol.number() in numbers + } + + with sire_lock: + system0 = _sr.load(str(topology0), show_warnings=False) + system1 = _sr.load(str(topology1), show_warnings=False) + mols0 = system0.molecules() + molecules0 = list(mols0) + molecules1 = list(system1.molecules()) + + starts = [] + start = 0 + for mol in molecules0: + starts.append(start) + start += mol.num_atoms() + + def unchanged(i): + mol0, mol1 = molecules0[i], molecules1[i] + charges0 = [q.value() for q in mol0.property("charge").to_list()] + charges1 = [q.value() for q in mol1.property("charge").to_list()] + types0 = mol0.property("ambertype").to_list() + return ( + charges0 == charges1 and types0 == mol1.property("ambertype").to_list() + ) + + # GCMC leaves out perturbable waters, e.g. co-alchemical ions, which + # are only water at one end state. + indices = [i for i in sorted(waters(system0) & waters(system1)) if unchanged(i)] + + by_molecule = {} + if indices and selection is not None: + try: + atoms = system0[selection].atoms() + except Exception: + atoms = [] + numbers = {mol.number(): i for i, mol in enumerate(mols0)} + for atom in atoms: + index = numbers[atom.molecule().number()] + by_molecule.setdefault(index, []).append(atom.index().value()) + + layout = None + if indices: + water = mols0[indices[0]] + elements = [a.element().symbol() for a in water.atoms()] + layout = { + "waters": _np.array([starts[i] for i in indices]), + "oxygen": elements.index("O"), + # Any massless site, e.g. of TIP4P, isn't shown. + "atoms": [ + (i, a.name().value(), e) + for i, (a, e) in enumerate(zip(water.atoms(), elements)) + if e in ("O", "H") + ], + "selection": [ + ({"first": starts[i], "num_atoms": mols0[i].num_atoms()}, atoms) + for i, atoms in sorted(by_molecule.items()) + ], + } + + with _topologies_lock: + _water_layouts[key] = layout + while len(_water_layouts) > _MAX_TOPOLOGIES: + _water_layouts.popitem(last=False) + return layout + + def _atoms(mol, protein): """ The atoms of an end state of a molecule that are shown: their indices in diff --git a/src/somd2/_viewer/_conformers.py b/src/somd2/_viewer/_conformers.py index 4c86f0e9..fef69844 100644 --- a/src/somd2/_viewer/_conformers.py +++ b/src/somd2/_viewer/_conformers.py @@ -112,8 +112,11 @@ def coordinates(p=p, n=n): mol0, stereo0 = _with_stereo(p.rdmol0, p.map0, coordinates) mol1, stereo1 = _with_stereo(p.rdmol1, p.map1, coordinates) - _embed(mol0) - _embed(mol1, mol0, p.mapping) + # An end state needn't have any atoms, e.g. of an absolute run. + if mol0.GetNumAtoms(): + _embed(mol0) + if mol1.GetNumAtoms(): + _embed(mol1, mol0 if mol0.GetNumAtoms() else None, p.mapping) index0 = {merged: i for i, merged in enumerate(p.map0)} index1 = {merged: i for i, merged in enumerate(p.map1)} diff --git a/src/somd2/_viewer/_coordinates.py b/src/somd2/_viewer/_coordinates.py index f0edd846..841cf701 100644 --- a/src/somd2/_viewer/_coordinates.py +++ b/src/somd2/_viewer/_coordinates.py @@ -47,7 +47,13 @@ class SavedCoordinates: """ def __init__( - self, positions, box=None, time_ps=None, virtual_sites=None, atoms_only=False + self, + positions, + box=None, + time_ps=None, + virtual_sites=None, + atoms_only=False, + ghost_waters=None, ): self.positions = positions self.box = box @@ -56,6 +62,9 @@ def __init__( # unknown. Not needed if the coordinates are of the atoms alone. self.virtual_sites = virtual_sites self.atoms_only = atoms_only + # The indices of the GCMC waters that are ghosts, among all of the + # waters that GCMC can move, or None if unknown. + self.ghost_waters = ghost_waters def molecule(self, first_atom, num_atoms, num_system_atoms): """ @@ -79,6 +88,25 @@ def molecule(self, first_atom, num_atoms, num_system_atoms): raise Mismatch() return self.positions[start : start + num_atoms] + def molecule_starts(self, first_atoms, num_system_atoms): + """ + The indices among the coordinates of the first atoms of molecules, + given the indices of those atoms in a topology of num_system_atoms + atoms. + """ + first_atoms = _np.asarray(first_atoms, dtype=int) + if self.atoms_only: + return first_atoms + if self.virtual_sites is not None and len(self.virtual_sites): + # Shifted by the virtual sites of the molecules before each. + rows = _np.asarray(self.virtual_sites, dtype=int) + rows = rows[_np.argsort(rows[:, 0])] + shifts = _np.concatenate([[0], _np.cumsum(rows[:, 1])]) + return first_atoms + shifts[_np.searchsorted(rows[:, 0], first_atoms)] + if self.virtual_sites is not None or len(self.positions) == num_system_atoms: + return first_atoms + raise Mismatch() + def read_coordinates(path, window=0): """ @@ -97,6 +125,7 @@ def read_coordinates(path, window=0): box=checkpoint.get("box"), time_ps=float(time_ps[0]) if time_ps is not None else None, virtual_sites=checkpoint.get("virtual_sites"), + ghost_waters=checkpoint.get("ghost_water_indices"), ) if path.suffix == ".s3": @@ -136,6 +165,9 @@ def read_repex_coordinates(path, windows): if not hasattr(state, "_num_slots"): order = _np.argsort(state._states) windows = [int(order[window]) for window in windows] + # Stored in the same order as the OpenMM states, with 1 for a water that is + # present and 0 for a ghost. + water_states = getattr(state, "_gcmc_states", None) or [None] * len(states) def nanometres(value): return _np.asarray(value.value_in_unit(_unit.nanometer)) @@ -158,6 +190,11 @@ def nanometres(value): box=nanometres(box) if box is not None else None, time_ps=_picoseconds(getattr(state, "_time", None)), virtual_sites=getattr(state, "_virtual_sites", None), + ghost_waters=( + None + if water_states[window] is None + else _np.where(_np.asarray(water_states[window]) == 0)[0] + ), ) ) return result diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index b1188820..3984e602 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -94,6 +94,11 @@ def _run(self): _block_re = _re.compile(r"^Finished block (\d+) of (\d+) for \S+ = ([\d.]+)") _perf_re = _re.compile(r"^Overall performance: ([\d.]+) ns day-1") _correction_re = _re.compile(r"Standard state correction: (-?[\d.]+) kcal mol-1") +# Older versions don't log the time. +_waters_re = _re.compile( + r"^Current number of waters in GCMC volume at \S+ = ([\d.]+)" + r"(?: \(time = ([\d.]+) ns\))? is (\d+)" +) def discover(paths, strict=True): @@ -248,6 +253,9 @@ def _reset(self): self._saved_states = 0 self.standard_state_correction = None self.messages = _deque(maxlen=20) + # The number of waters in the GCMC region at each checkpoint, by λ, as + # (time in ns, or None if not logged, count). + self.gcmc_waters = {} self._new_session(None) def _new_session(self, timestamp): @@ -326,6 +334,18 @@ def _parse(self, line): self.finished = True elif m := _correction_re.search(message): self.standard_state_correction = float(m.group(1)) + elif m := _waters_re.match(message): + time = float(m.group(2)) if m.group(2) is not None else None + counts = self.gcmc_waters.setdefault(float(m.group(1)), []) + # A block run again after a restart replaces its earlier count. + while ( + time is not None + and counts + and counts[-1][0] is not None + and counts[-1][0] >= time + ): + counts.pop() + counts.append((time, int(m.group(3)))) if level in ("WARNING", "ERROR", "CRITICAL"): self.messages.append((_parse_timestamp(stamp), level, message)) @@ -407,6 +427,26 @@ def load(): return self._cached("config", _stamp(path), load) + def _config_values(self, key): + """ + The values of an option in every config written for the run, e.g. by + restarts, with None where it wasn't recorded. + """ + from ..io import yaml_to_dict + + paths = sorted(self.path.glob("config*.yaml")) + + def load(): + values = set() + for path in paths: + try: + values.add(yaml_to_dict(str(path)).get(key)) + except Exception: + values.add(None) + return values + + return self._cached(f"config_values:{key}", _stamp(*paths), load) + def _lambdas(self, config): """ Return the sampled and energy lambda values, as set up by the runner. @@ -943,7 +983,7 @@ def _replica_paths(states, num_states, max_points=500): "round_trips": round_trips, } - def _samplers(self, config, lambda_values, repex_state): + def _samplers(self, config, lambda_values, repex_state, log=None): gcmc = {} flips = {} @@ -987,17 +1027,50 @@ def load(path=path): result = {} if config.get("gcmc") and gcmc: + # The statistics cover the whole run, so this is only known if it + # was never changed on restart, and was recorded, which it is by + # versions that let it be set. + attempts = config.get("gcmc_num_attempts") + if len(self._config_values("gcmc_num_attempts")) > 1: + attempts = None rows = [] for lam in sorted(gcmc): - row = {"lambda": lam} - row.update( - {k: v for k, v in gcmc[lam].items() if isinstance(v, (int, float))} - ) - moves = row.get("num_moves") + stats = gcmc[lam] + moves = stats.get("num_moves") + row = { + "lambda": lam, + "moves": moves, + "insertions": stats.get("num_insertions"), + "deletions": stats.get("num_deletions"), + "accepted_per_move": None, + "attempt_acceptance": None, + } if moves: - row["acceptance"] = row.get("num_accepted", 0) / moves + row["accepted_per_move"] = stats.get("num_accepted", 0) / moves + if attempts and "num_accepted_attempts" in stats: + row["attempt_acceptance"] = stats["num_accepted_attempts"] / ( + moves * attempts + ) rows.append(row) result["gcmc"] = rows + result["gcmc_batch_size"] = config.get("gcmc_batch_size") + result["gcmc_num_attempts"] = attempts + if config.get("gcmc") and log is not None and log.gcmc_waters: + # Without logged times, counts are only known by checkpoint. + timed = all(t is not None for c in log.gcmc_waters.values() for t, _ in c) + result["gcmc_waters"] = { + "timed": timed, + "windows": [ + { + "lambda": lam, + "time": [ + t if timed else i + 1 for i, (t, _) in enumerate(counts) + ], + "count": [n for _, n in counts], + } + for lam, counts in sorted(log.gcmc_waters.items()) + ], + } if config.get("terminal_flip_frequency") and flips: result["terminal_flip"] = [ { @@ -1251,7 +1324,13 @@ def binding_site(self, known=None): last = self._positions_file(last=True) if last.suffix != ".pkl" and _checkpoint_lambda(last) != 1.0: last = None - stamp = (_stamp(top0, top1), _stamp(path), _stamp(last) if last else None) + gcmc = self._gcmc_region() + stamp = ( + _stamp(top0, top1), + _stamp(path), + _stamp(last) if last else None, + repr(gcmc), + ) key = _hashlib.sha1(repr(stamp).encode()).hexdigest()[:16] if known == key: return {"status": "unchanged", "key": key} @@ -1269,7 +1348,7 @@ def load(): "status": "unavailable", "reason": "No coordinates have been saved yet.", } - site = binding_site(top0, top1, saved, saved1) + site = binding_site(top0, top1, saved, saved1, gcmc) except Mismatch: return { "status": "unavailable", @@ -1294,6 +1373,33 @@ def load(): # So that the page can title the card before it is shown. return _clean(dict(result, kind=kind)) + def _gcmc_region(self): + """ + With GCMC, the selection of atoms at the centre of its region and the + region's radius in nm, which are None if it samples the whole box. + None without GCMC. + """ + path = _find_config(self.path) + if path is None: + return None + + def load(): + import sire as _sr + + config = self.config() or {} + if not config.get("gcmc"): + return None + selection = config.get("gcmc_selection") + try: + radius = _sr.u(str(config["gcmc_radius"])).to("nm") + except Exception: + radius = None + if not selection or radius is None: + return {"selection": None, "radius": None} + return {"selection": str(selection), "radius": radius} + + return self._cached("gcmc_region", _stamp(path), load) + def _positions_file(self, last=False): """ The file holding the system's coordinates in the first, or last, λ @@ -1616,7 +1722,7 @@ def _summary(self): ), "samplers": self._section( "summarise the Monte Carlo samplers", - lambda: self._samplers(config, lambda_values, repex_state), + lambda: self._samplers(config, lambda_values, repex_state, log), ), "schedule": self._section( "summarise the λ schedule", diff --git a/src/somd2/_viewer/_depict.py b/src/somd2/_viewer/_depict.py index 3402d49c..2b550bfe 100644 --- a/src/somd2/_viewer/_depict.py +++ b/src/somd2/_viewer/_depict.py @@ -125,11 +125,17 @@ def depict(topology0, topology1): "smiles1": Chem.MolToSmiles(Chem.RemoveHs(p.rdmol1, sanitize=False)), "pixels": pixels, "mapping_pixels": mapping_pixels, - "svg_mapping": _draw_mapping( - p.rdmol0, p.rdmol1, p.mapping, p.map0, p.map1, mapping_pixels + # Only drawn for end states with atoms, which an absolute run + # needn't have at λ = 1. + "svg_mapping": ( + _draw_mapping( + p.rdmol0, p.rdmol1, p.mapping, p.map0, p.map1, mapping_pixels + ) + if p.rdmol0.GetNumAtoms() and p.rdmol1.GetNumAtoms() + else None ), - "svg0": _draw(p.rdmol0, pixels), - "svg1": _draw(p.rdmol1, pixels), + "svg0": _draw(p.rdmol0, pixels) if p.rdmol0.GetNumAtoms() else None, + "svg1": _draw(p.rdmol1, pixels) if p.rdmol1.GetNumAtoms() else None, } ) @@ -445,9 +451,14 @@ def _to_rdkit(mol, dummy): Returns the RDKit molecule, and the merged molecule index of each atom. """ import sire as _sr + from rdkit import Chem real = [i for i, d in enumerate(dummy) if not d] + # E.g. a molecule decoupled by turning all of its atoms into dummies. + if not real: + return Chem.Mol(), real + # Drop the parameters, which can't be carried over to a subset of atoms. mol = mol.edit().remove_property("parameters").commit() if len(real) < len(dummy): diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index b348cda0..1dd27bac 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -316,7 +316,7 @@

SOMD2 Viewer

try { return new Set(JSON.parse(localStorage.getItem(key) || "[]")); } catch (e) { return new Set(); } } -const state = { view: ["#summary", "#network"].includes(location.hash) ? location.hash.slice(1) : "run", summaryData: null, summaryAvailable: false, summaryTimer: null, networkData: null, networkAvailable: false, runs: [], current: null, summary: null, depictions: {}, timer: null, quickTimer: null, depictionsLoading: {}, collapsed: loadCollapsed(), navCollapsed: loadCollapsed("somd2-viewer-nav-collapsed"), components: {}, componentsLambda: {}, configFilter: "", problems: {}, requests: 0, latest: {}, shown: {}, group: null, repeatsOpen: new Map(), conformers: {}, conformersLoading: {}, structureMode: "0", structureMolecule: {}, viewer3d: null, bindingSites: {}, bindingSitesLoading: {}, viewerSite: null, siteMode: "0" }; +const state = { view: ["#summary", "#network"].includes(location.hash) ? location.hash.slice(1) : "run", summaryData: null, summaryAvailable: false, summaryTimer: null, networkData: null, networkAvailable: false, runs: [], current: null, summary: null, depictions: {}, timer: null, quickTimer: null, depictionsLoading: {}, collapsed: loadCollapsed(), navCollapsed: loadCollapsed("somd2-viewer-nav-collapsed"), components: {}, componentsLambda: {}, configFilter: "", problems: {}, requests: 0, latest: {}, shown: {}, group: null, repeatsOpen: new Map(), conformers: {}, conformersLoading: {}, structureMode: "0", structureMolecule: {}, viewer3d: null, bindingSites: {}, bindingSitesLoading: {}, viewerSite: null, siteMode: "0", siteWaters: true, gcmcWatersLambda: {} }; const PALETTE = ["#2f6fdf", "#e0663b", "#2a9d6f", "#b04fc9", "#c9a227", "#3aa7c9", "#d64f7a", "#7a8699"]; const LAMBDA = "λ"; // The thresholds for poor overlap and replica mixing, as in the summary. @@ -955,7 +955,8 @@

SOMD2 Viewer

let html = ""; for (const [i, m] of d.molecules.entries()) { html += `

${esc(m.name)}

`; - if (absolute) { + // Without a mapping, e.g. if an end state has no atoms, only λ = 0 is drawn. + if (absolute || !m.svg_mapping) { const verb = s.schedule.name === "annihilate" ? "Annihilated" : "Decoupled"; const width = (m.pixels || 300) + 18; html += `
${m.svg0}
${verb} at ${LAMBDA} = 1
${esc(m.smiles0)}
${SAVE_BUTTONS}
`; @@ -1173,6 +1174,7 @@

SOMD2 Viewer

const modes = [["0", `${LAMBDA} = 0`], ["1", `${LAMBDA} = 1`]]; controls += modes.map(([value, label]) => ``).join(""); } + controls += ``; controls += ``; if (document.fullscreenEnabled) controls += fullScreenButton(); let html = `
${controls}
`; @@ -1187,9 +1189,21 @@

SOMD2 Viewer

} else { html += `
The latest saved coordinates of the ${LAMBDA} = 0 window${time}, updated as the simulation runs. Residues within 6 Å of the perturbed molecule are shown in detail, and the rest of the protein as a cartoon. Drag to rotate, Ctrl + scroll to zoom.
`; } + html += waterNote(b); return card(title, html); } +// Which waters are shown, and how many. +function waterNote(b) { + const mutation = b.kind === "mutation"; + const states = mutation ? b.states : [b]; + const counts = states.map((s, i) => `${s.num_waters}${mutation ? ` at ${LAMBDA} = ${i}` : ""}`).join(" and "); + const near = mutation ? "the mutated residues" : "the perturbed molecule"; + const region = states.map((s) => s.gcmc).find(Boolean); + if (!region) return `
Waters within 6 Å of ${near} are shown (${counts}).
`; + return `
Waters within 6 Å of ${near}, or in the GCMC region, are shown (${counts}), leaving out GCMC's ghost waters. The region, within ${fmt(region.radius, 1)} Å of the centre of ${esc(region.selection)}, is shown as a wireframe sphere.
`; +} + // Polled while the card is open, sending the version already shown so that an // unchanged structure isn't sent again. function loadBindingSite(id) { @@ -1256,19 +1270,33 @@

SOMD2 Viewer

v.site = site; v.viewer.removeAllModels(); if (site.kind === "mutation") { - v.models = site.states.map((s) => v.viewer.addModel(s.pdb, "pdb")); + v.models = site.states.map((s) => v.viewer.addModel(s.pdb, "pdb", { keepH: true })); styleSite(v); return; } - const model = v.viewer.addModel(site.pdb, "pdb"); + // Hydrogens are kept for the waters, but only they have them drawn. + const model = v.viewer.addModel(site.pdb, "pdb", { keepH: true }); model.setStyle({}, { cartoon: { color: "lightgrey" } }); // Added to the cartoon, rather than replacing it. styleSideChains(model, site.pocket, { stick: { radius: 0.15, colorscheme: "grayCarbon" } }); const ligand = { colorscheme: "greenCarbon" }; model.setStyle({ serial: site.ligand }, { stick: { radius: 0.14, ...ligand }, sphere: { scale: 0.22, ...ligand } }); + model.setStyle({ serial: site.ligand, elem: "H" }, {}); + styleWaters(v, model, site.waters, site.gcmc); v.viewer.render(); } +// The waters, in ball-and-stick, and any GCMC region as a wireframe sphere, +// unless the waters are turned off. +function styleWaters(v, model, waters, gcmc) { + v.viewer.removeAllShapes(); + if (!state.siteWaters) return; + model.setStyle({ serial: waters }, { stick: { radius: 0.1 }, sphere: { scale: 0.2 } }); + if (!gcmc) return; + const [x, y, z] = gcmc.centre; + v.viewer.addSphere({ center: { x, y, z }, radius: gcmc.radius, color: css("--accent"), wireframe: true, opacity: 0.3 }); +} + // Draw residues over the cartoon from Cα outwards, so that their backbone // doesn't overlap it. A proline keeps its N to close its ring. const BACKBONE = new Set(["N", "C", "O", "OXT"]); @@ -1292,6 +1320,8 @@

SOMD2 Viewer

model.setStyle({ serial: s.mutated }, { stick: { radius: 0.18, colorscheme: "magentaCarbon" } }, true); const other = { colorscheme: "greenCarbon" }; model.setStyle({ serial: s.ligand }, { stick: { radius: 0.14, ...other }, sphere: { scale: 0.22, ...other } }); + model.setStyle({ serial: s.ligand, elem: "H" }, {}); + styleWaters(v, model, s.waters, s.gcmc); }); v.viewer.render(); } @@ -1402,10 +1432,24 @@

SOMD2 Viewer

if (sm.error) return card("Monte Carlo samplers", `
Couldn't summarise the samplers: ${esc(sm.error)}
`); let html = ""; if (sm.gcmc) { - const keys = Object.keys(sm.gcmc[0]).filter((k) => k !== "lambda"); - html += `

GCMC

${keys.map((k) => ``).join("")}`; - for (const row of sm.gcmc) html += `${keys.map((k) => ``).join("")}`; + const perMove = (value) => { + // More than one acceptance per move wastes the rest of each batch. + const wasteful = value > 1; + const style = wasteful ? ' style="color:var(--warn)" title="More than one accepted per move, so trials are wasted"' : ""; + return ``; + }; + html += `

GCMC

${LAMBDA}${esc(k.replace(/^num_/, "").replace(/_/g, " "))}
${fmtLambda(row.lambda)}${k === "acceptance" ? fmt(row[k], 3) : row[k] ?? "–"}
${fmt(value, 2)}
`; + for (const row of sm.gcmc) html += `${perMove(row.accepted_per_move)}`; html += `
${LAMBDA}movesinsertionsdeletionsaccepted per moveacceptance per attempt
${fmtLambda(row.lambda)}${row.moves ?? "–"}${row.insertions ?? "–"}${row.deletions ?? "–"}${fmt(row.attempt_acceptance, 4)}
`; + const settings = sm.gcmc_num_attempts ? ` This run used a batch size of ${sm.gcmc_batch_size} and ${sm.gcmc_num_attempts} attempts per move.` : ""; + html += `
Each GCMC move runs its trials in parallel batches, so it can accept more than one insertion or deletion. More than one accepted per move, highlighted, means trials are wasted, so gcmc_batch_size and gcmc_num_attempts should be reduced, aiming for about one per move. The acceptance per attempt is for the approximate reaction-field potential used to choose candidates.${settings}
`; + if (sm.gcmc_waters) { + const shown = state.gcmcWatersLambda[s.id] ?? "all"; + const options = [``, ...sm.gcmc_waters.windows.map((w) => ``)].join(""); + html += `

Waters in the GCMC region

`; + const timing = sm.gcmc_waters.timed ? "" : " This run's log doesn't record the time of each count, so they're shown by checkpoint."; + html += `
At each checkpoint, from the log.${timing}
`; + } } if (sm.terminal_flip) { html += `Terminal flips
`; @@ -1586,6 +1630,24 @@

SOMD2 Viewer

xLabel: LAMBDA, yLabel: "scale factor", xMin: 0, xMax: 1, }); } + const gw = $("#gcmc-waters-chart"); + const gcmcWaters = s.samplers && s.samplers.gcmc_waters; + if (gw && gcmcWaters) { + const allWaters = gcmcWaters.windows; + const timed = gcmcWaters.timed; + const shown = state.gcmcWatersLambda[s.id] ?? "all"; + // Each window keeps its colour from the full set. + const waters = allWaters.map((w, i) => ({ ...w, colour: colour(i, allWaters.length) })).filter((w) => shown === "all" || String(w.lambda) === String(shown)); + lineChart(gw, { + series: waters.map((w) => ({ + x: w.time, y: w.count, color: w.colour, name: `${LAMBDA} = ${fmtLambda(w.lambda)}`, + points: w.count.length <= 50, title: (x, y) => `${LAMBDA} = ${fmtLambda(w.lambda)}, ${timed ? `${fmt(x, 4)} ns` : `checkpoint ${x}`}: ${y} waters`, + })), + xLabel: timed ? "time (ns)" : "checkpoint", yLabel: "waters in region", yMin: 0, + yFormat: (v) => (Number.isInteger(v) ? String(v) : ""), + legend: waters.length <= 16, + }); + } } function valueWithError(value, error) { @@ -2310,6 +2372,14 @@

SOMD2 Viewer

const v = state.viewerSite; if (v && v.viewer && v.site && v.site.kind === "mutation") styleSite(v); })); + const siteWaters = $("#site-waters"); + if (siteWaters) siteWaters.addEventListener("change", (e) => { + state.siteWaters = e.target.checked; + const v = state.viewerSite; + if (!v || !v.viewer || !v.site) return; + if (v.site.kind === "mutation") styleSite(v); + else showSite(v, v.site); + }); const resetSite = $("#binding-site-reset"); if (resetSite) resetSite.addEventListener("click", () => { if (state.viewerSite && state.viewerSite.viewer) fitStructure(state.viewerSite); @@ -2327,6 +2397,11 @@

SOMD2 Viewer

state.structureMolecule[s.id] = Number(e.target.value); render(); }); + const watersSelect = $("#gcmc-waters-lambda"); + if (watersSelect) watersSelect.addEventListener("change", (e) => { + state.gcmcWatersLambda[s.id] = e.target.value; + renderCharts(s); + }); const lambdaSelect = $("#components-lambda"); if (lambdaSelect) lambdaSelect.addEventListener("change", (e) => { state.componentsLambda[s.id] = parseFloat(e.target.value); diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 33fa0684..da9b69d9 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -119,6 +119,8 @@ class Config: "max_contexts", "restraint_search_time", "restraint_search_frequency", + "gcmc_batch_size", + "gcmc_num_attempts", ] ) @@ -199,6 +201,8 @@ def __init__( gcmc_radius="4 A", gcmc_bulk_sampling_probability=0.1, gcmc_tolerance=0.0, + gcmc_batch_size=1000, + gcmc_num_attempts=10000, use_dispersion_correction=False, rest2_scale=1.0, rest2_selection=None, @@ -558,6 +562,15 @@ def __init__( of acceptance for a move. This can be used to exclude low probability candidates that can cause instabilities or crashes for the MD engine. + gcmc_batch_size: int + The number of GCMC insertion and deletion trials evaluated in parallel in each + batch. This should be tuned so that an average of about one trial is accepted + per batch, since the rest of the batch is wasted. + + gcmc_num_attempts: int + The total number of GCMC trials per move, which must be at least + 'gcmc_batch_size'. + use_dispersion_correction: bool Whether to use the long-range dispersion correction for LJ interactions. When True, the correction is evaluated analytically via a CustomVolumeForce @@ -786,6 +799,8 @@ def __init__( self.gcmc_radius = gcmc_radius self.gcmc_bulk_sampling_probability = gcmc_bulk_sampling_probability self.gcmc_tolerance = gcmc_tolerance + self.gcmc_batch_size = gcmc_batch_size + self.gcmc_num_attempts = gcmc_num_attempts self.use_dispersion_correction = use_dispersion_correction self.rest2_scale = rest2_scale self.rest2_selection = rest2_selection @@ -2413,6 +2428,46 @@ def gcmc_tolerance(self, gcmc_tolerance): raise ValueError("'gcmc_tolerance' must be greater than or equal to 0.0") self._gcmc_tolerance = gcmc_tolerance + @property + def gcmc_batch_size(self): + return self._gcmc_batch_size + + @gcmc_batch_size.setter + def gcmc_batch_size(self, gcmc_batch_size): + if not isinstance(gcmc_batch_size, int): + try: + gcmc_batch_size = int(gcmc_batch_size) + except Exception: + raise ValueError("'gcmc_batch_size' must be an integer") + if gcmc_batch_size <= 0: + raise ValueError("'gcmc_batch_size' must be greater than 0") + # Only set after the batch size on construction. + if ( + hasattr(self, "_gcmc_num_attempts") + and gcmc_batch_size > self._gcmc_num_attempts + ): + raise ValueError( + "'gcmc_num_attempts' must be greater than or equal to 'gcmc_batch_size'" + ) + self._gcmc_batch_size = gcmc_batch_size + + @property + def gcmc_num_attempts(self): + return self._gcmc_num_attempts + + @gcmc_num_attempts.setter + def gcmc_num_attempts(self, gcmc_num_attempts): + if not isinstance(gcmc_num_attempts, int): + try: + gcmc_num_attempts = int(gcmc_num_attempts) + except Exception: + raise ValueError("'gcmc_num_attempts' must be an integer") + if gcmc_num_attempts < self._gcmc_batch_size: + raise ValueError( + "'gcmc_num_attempts' must be greater than or equal to 'gcmc_batch_size'" + ) + self._gcmc_num_attempts = gcmc_num_attempts + @property def use_dispersion_correction(self): return self._use_dispersion_correction diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 986627d3..74d24606 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1021,6 +1021,8 @@ def __init__(self, system, config): "taylor_power": self._config.taylor_power, "standard_volume": str(self._config.gcmc_standard_volume), "tolerance": self._config.gcmc_tolerance, + "batch_size": self._config.gcmc_batch_size, + "num_attempts": self._config.gcmc_num_attempts, } else: self._gcmc_kwargs = None diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index df82a301..4604947b 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2883,7 +2883,7 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): ) _logger.info( f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} " - f"is {num_waters}{acc_str}" + f"(time = {system.time().to('ns'):.4f} ns) is {num_waters}{acc_str}" ) # Log terminal flip acceptance rate for this replica. diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index a54c9768..bae8ec32 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -1176,7 +1176,7 @@ def generate_lam_vals(lambda_base, increment=0.001): ) _logger.info( f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} " - f"is {num_waters}{acc_str}" + f"(time = {system.time().to('ns'):.4f} ns) is {num_waters}{acc_str}" ) finally: gcmc_sampler.pop()
${LAMBDA}attemptedacceptedacceptance