"""NucDynamics on the native engine (``uchrom_nucdyn``, Rust; CPU and GPU).
The engine is a separate package (``pip install u-chrom[nucdyn]`` or
``pip install ./native/nucdyn``); this module turns contacts into a
:class:`~uchrom.core.ChromData` ensemble.
"""
from __future__ import annotations
import os
import warnings
from typing import Any, Dict, Mapping, Optional, Sequence, Union
import numpy as np
import pandas as pd
__all__ = ["reconstruct_nucdyn", "native_available", "engine_info", "load_contacts"]
#: Default hierarchical schedule (bp), the original's ``-s 8 4 2 0.4 0.2 0.1``.
DEFAULT_PARTICLE_SIZES = (8e6, 4e6, 2e6, 4e5, 2e5, 1e5)
#: Public names of engine protocols: EMber is protocol "robust" of the native engine.
PROTOCOL_ALIASES = {"ember": "robust"}
def _engine():
import uchrom_nucdyn
return uchrom_nucdyn
[docs]
def native_available() -> bool:
"""True when the native engine (``uchrom_nucdyn``) can be imported."""
try:
_engine()
except ImportError:
return False
return True
[docs]
def engine_info() -> Dict[str, Any]:
"""Version of the native engine and the GPU it would use."""
un = _engine()
return {"version": un.__version__, "gpu_build": bool(un.HAS_GPU), "gpu": un.gpu_info()}
def _resolve_device(device: str) -> str:
if device in ("cpu", "gpu"):
return device
if device != "auto":
raise ValueError(f"device must be 'auto', 'cpu' or 'gpu', got {device!r}")
return "gpu" if _engine().gpu_available() else "cpu"
[docs]
def reconstruct_nucdyn(
contacts: Union[str, os.PathLike, pd.DataFrame],
*,
n_models: int = 10,
device: str = "auto",
n_threads: int = 0,
seed: int = 0,
protocol: str = "nuc_dynamics_2017",
particle_sizes: Sequence[float] = DEFAULT_PARTICLE_SIZES,
genome_ranges=None,
cell_id: Optional[str] = None,
key_added: Optional[str] = "nucdyn",
params: Optional[Mapping[str, Any]] = None,
verbose: bool = False,
_method: str = "nucdyn",
**engine_params: Any,
):
"""Single-cell Hi-C genome structure ensemble with NucDynamics.
Runs the hierarchical simulated-annealing protocol of NucDynamics
(Stevens et al. 2017, *Nature* 544:59) on the native engine and returns
the ensemble as a new :class:`~uchrom.core.ChromData`.
Parameters
----------
contacts
Path (``.ncc``, ``.pairs[.gz]``, GEO ``chr_A pos_A chr_B pos_B``
table) or DataFrame of contacts, see :func:`load_contacts`.
n_models
Number of structures (models) in the ensemble. All models are
computed at once (in parallel on the CPU, in one batch on the GPU).
device
``"auto"`` (GPU when available, else CPU), ``"cpu"`` (f64, multi-core)
or ``"gpu"`` (f32, wgpu: Metal / Vulkan / DX12).
n_threads
CPU worker threads (0 = all cores).
seed
Random seed (start coordinates and velocities); on the CPU a seed
gives bitwise-identical results.
protocol
``"nuc_dynamics_2017"`` (default): the code that produced the
published 2017 structures (tjs23/nuc_dynamics ``master``).
``"release_1.3"``: the later bead-size-scaled version with
ambiguity resolution and model selection (2 x ``n_models`` above 1 Mb,
the ``n_models`` closest to the mean kept).
``"robust"`` (public name **EMber**, alias ``"ember"``; see
:func:`uchrom.recon.sc.ember.reconstruct_ember`): the 2017 protocol with a noise-aware observation model
(contact restraints re-weighted by their posterior probability of
being true contacts; false-contact rate and contact-kernel width
learned by EM; ``robust_*`` engine parameters). Adds the results
``<key_added>.em`` (EM log) and ``<key_added>.contact_weights``
(final weight of every input contact, NaN = filtered out).
particle_sizes
Hierarchical particle sizes in bp (coarse to fine).
genome_ranges
Restrict the contacts to these ranges (e.g. ``"chr19"``).
cell_id
Written to ``spots['cell_id']``.
key_added
Key of the provenance record in ``cd.results`` (per-stage log) and
of ``cd.uns[key_added]``; ``None`` stores neither.
params
Further engine parameters (see ``uchrom_nucdyn.default_params()``:
``temp_steps``, ``dynamics_steps``, ``temp_start``, ``temp_end``,
``time_step``, ``contact_dist_lower``, ...); keyword arguments are
merged into it.
Returns
-------
ChromData
One spot per particle; one trace per chromosome. ``coords`` holds
model 0 and ``layers['model_<k>']`` holds every model ``k`` (model 0
included), so the whole ensemble travels with the object. Units are
particle radii. Bins are the particles' genomic intervals: in the
2017 protocol a particle at sequence position ``p`` collects the
contacts in ``(p - size, p]`` (``bins.start = p - size``,
``bins.end = p``); in release_1.3 ``[p, p + size)``. ``uns[key_added]``
records the engine, protocol, parameters and per-stage log.
"""
from uchrom.core import ChromData
un = _engine()
protocol = PROTOCOL_ALIASES.get(protocol, protocol)
func = f"uchrom.recon.sc.{_method}.reconstruct_{_method}"
c = load_contacts(contacts, genome_ranges)
dev = _resolve_device(device)
p = dict(params or {})
p.update(engine_params)
res = un.calc_genome_structure(
c["chrom_a"], c["pos_a"], c["chrom_b"], c["pos_b"],
ambiguity=c["ambiguity"], active=c["active"],
n_models=int(n_models), device=dev, n_threads=int(n_threads), seed=int(seed),
protocol=protocol, particle_sizes=[float(s) for s in particle_sizes], verbose=verbose, **p)
size = int(res["particle_size"])
pos = np.asarray(res["position"], np.int64)
if res["position_convention"] == "end":
start = np.maximum(pos - size, 0)
end = np.maximum(pos, start + 1)
else:
start, end = pos, pos + size
coords = np.asarray(res["coords"], np.float64)
df = pd.DataFrame({"chrom": res["chrom"].astype(str), "start": start, "end": end,
"x": coords[0, :, 0], "y": coords[0, :, 1], "z": coords[0, :, 2]})
cd = ChromData.from_dataframe(df, cell_id=cell_id)
# from_dataframe keeps the row order of df; every model in the same order
width = max(1, len(str(len(coords) - 1)))
for k in range(len(coords)):
cd.layers[f"model_{k:0{width}d}"] = coords[k].copy()
cd.uns.setdefault("xyz_unit", "particle radii (NucDynamics)")
if key_added is not None:
stages = pd.DataFrame([{k: v for k, v in s.items() if k not in ("removed_models", "robust")}
for s in res["stages"]])
info = {
"method": _method, "engine": "uchrom_nucdyn", "engine_version": un.__version__, "device": dev,
"gpu": un.gpu_info() if dev == "gpu" else None, "protocol": protocol,
"n_models": int(len(coords)), "seed": int(seed), "particle_size": size,
"particle_position": res["position_convention"],
"model_layers": [f"model_{k:0{width}d}" for k in range(len(coords))],
"seconds": float(res["seconds_total"]),
"citation": ("U-Chrom EMber (engine protocol \"robust\"), built on " if protocol == "robust" else "")
+ "NucDynamics: Stevens et al. 2017, Nature 544:59, doi:10.1038/nature21429",
}
cd.uns[key_added] = info
cd.results.set(f"{key_added}.stages", stages, kind="table",
function=func,
params={k: v for k, v in res["params"].items() if k != "verbose"},
inputs={"contacts": str(contacts) if not isinstance(contacts, pd.DataFrame) else "DataFrame",
"n_contacts": int(len(c["pos_a"]))})
if protocol == "robust":
# EM log (one row per update) and the final weight of every input contact
em = pd.DataFrame([{"stage": s["stage"], "particle_size": s["particle_size"], **u}
for s in res["stages"] for u in s.get("robust", [])])
cd.results.set(f"{key_added}.em", em, kind="table",
function=func, params={"protocol": "robust"})
cd.results.set(f"{key_added}.contact_weights", np.asarray(res["contact_weights"], np.float64),
kind="array", function=func,
params={"protocol": "robust"},
inputs={"rows": "input contacts in load_contacts order; NaN = filtered out"})
return cd