torchref.config

Centralized configuration for TorchRef.

Set at import time from the environment – TORCHREF_DTYPE_FLOAT (float32 default / float64), TORCHREF_DTYPE_INT (int32 / int64), TORCHREF_DTYPE_COMPLEX (complex64 / complex128), TORCHREF_DEVICE (‘auto’ default / ‘cuda’ / ‘mps’ / ‘cpu’), TORCHREF_SIGMA_CUTOFF_ED (3.0), TORCHREF_COMPILE_TARGETS and TORCHREF_CACHING (on by default) – or at runtime by attribute assignment:

torchref.dtypes.float = torch.float64
torchref.device.current = torch.device('cpu')
torchref.sigma_cutoff_ed.value = 4.0   # density splat truncation, in sigmas
torchref.config.caching.value = False  # recompute every cached forward()

The default device is auto-detected cuda -> mps -> cpu, and a CUDA device is picked only if its compute capability is >= the minimum sm_* in this PyTorch build and its VRAM is >= _MIN_CUDA_VRAM_GB; otherwise auto-detection falls back with a warning naming the failing requirement. An explicit TORCHREF_DEVICE bypasses those gates but still fails fast if the backend is unavailable.

MPS supports neither float64 nor complex128. Resolving to MPS with float64 configured warns at import; set TORCHREF_DTYPE_FLOAT=float32 or TORCHREF_DEVICE=cpu.

Functions

canonical_device(dev)	Return dev with its default index filled in (None passes through).
get_caching_enabled()	Whether CachedForwardMixin should serve cached forward() results.
get_compile_targets()	Whether quadrature X-ray target kernels should be ``torch.compile``d.
get_complex_dtype()	Get the current default complex dtype.
get_default_device()	Get the current default device.
get_float_dtype()	Get the current default float dtype.
get_int_dtype()	Get the current default int dtype.
get_sigma_cutoff_ed()	Get the current density-splat sigma cutoff (number of sigmas).
normalize_device([dev])	Coerce a user-supplied device (or None) to a canonical device.

Classes

CachingConfig()	Whether torchref.utils.CachedForwardMixin serves cached results.
CompileTargetsConfig()	Whether to torch.compile the quadrature X-ray target kernels.
DeviceConfig()	The active device: device.current reads it, assignment sets it.
DtypeConfig()	Dtype configuration by attribute: dtypes.float, .int, .complex, readable and assignable (dtypes.float = torch.float64).
SigmaCutoffConfig()	Number of sigmas at which the per-atom electron-density Gaussian is truncated.

class torchref.config.DtypeConfig

Dtype configuration by attribute: dtypes.float, .int, .complex, readable and assignable (dtypes.float = torch.float64).

__init__()

property float: dtype

Get the current default float dtype.

property int: dtype

Get the current default int dtype.

property complex: dtype

Get the current default complex dtype.

torchref.config.get_float_dtype()

Get the current default float dtype.

torchref.config.get_int_dtype()

Get the current default int dtype.

torchref.config.get_complex_dtype()

Get the current default complex dtype.

class torchref.config.SigmaCutoffConfig

Number of sigmas at which the per-atom electron-density Gaussian is truncated.

sigma_cutoff_ed.value reads or sets it; initialised from TORCHREF_SIGMA_CUTOFF_ED (default 3.0). Must be positive.

__init__()

property value: float

Get the current sigma cutoff (number of sigmas).

torchref.config.get_sigma_cutoff_ed()

Get the current density-splat sigma cutoff (number of sigmas).

class torchref.config.CompileTargetsConfig

Whether to torch.compile the quadrature X-ray target kernels.

compile_targets.value reads or sets it; initialised from TORCHREF_COMPILE_TARGETS (“1”/”true”/”yes”/”on”). Applies to the full-form MLF target (--xray-mode ml_full), whose per-reflection fixed-node quadrature is bound by dispatch and memory traffic in eager mode, so fusing it is worth roughly an order of magnitude.

Off by default because of compile latency, which autograd dominates: compiling the backward costs ~2 minutes on the first call, so a short refinement of a small structure gets slower, not faster. It pays off for big datasets (the target scales with reflection count) and for long or repeated runs in one process (ensembles, collection/PanDDA refinements, interactive sessions) where the compile amortises.

Only the reflection-count dimension varies, so the kernels compile with dynamic=True and one compilation serves every dataset size, work/free subset and gathered tensor – no chunking or padding layer is needed. To cut latency point TORCHINDUCTOR_CACHE_DIR at node-local disk (never gpfs) so codegen is reused across processes; artifacts are ~22 MB.

Keep it off for float64 and gradient-verification work regardless: that path is the eager reference and is deliberately unfused.

__init__()

property value: bool

Whether quadrature target kernels are ``torch.compile``d.

torchref.config.get_compile_targets()

Whether quadrature X-ray target kernels should be ``torch.compile``d.

class torchref.config.CachingConfig

Whether torchref.utils.CachedForwardMixin serves cached results.

caching.value reads or sets it; initialised from TORCHREF_CACHING (“1”/”true”/”yes”/”on” vs “0”/”false”/”no”/”off”), on by default. Turning it off makes the mixin inert: every module call runs forward() again, so ModelFT, MixedTensor, RigidXYZTensor and their subclasses recompute structure factors and parameter transforms from scratch on each access. The numbers are unchanged – only the work done to get them.

Gates that mixin only. The other hand-rolled caches (Cell._cache, SigmaAEstimator._cache, the bulk-solvent and parity caches, the reflection-data subset views, the symmetry-extractor rebuild in sf_fft) are unaffected.

Intended for diagnosis rather than production: if refinement produces stale-looking numbers, rerunning with TORCHREF_CACHING=0 says in one step whether the forward cache is responsible. Also useful as an eager reference when changing the mixin’s fingerprinting. Use torchref.utils.no_caching() to scope the change to a block.

__init__()

property value: bool

Whether cached forward() results are served.

torchref.config.get_caching_enabled()

Whether CachedForwardMixin should serve cached forward() results.

torchref.config.canonical_device(dev)

Return dev with its default index filled in (None passes through).

torch.device('cuda') != torch.device('cuda:0') although both name one physical device, and a device read off a real tensor always carries an index – so without a shared normal form obj.device == tensor.device is False on a fresh object and True after its first .to(). cpu deliberately stays bare, since a CPU tensor’s device has no index and cpu:0 would recreate the mismatch.

This is the allocation-free form of torch.empty(0, device=d).device, which matters on constructor and comparison paths. Deliberately not memoised: torch.cuda.current_device() is mutable via set_device, so a cache would freeze a stale answer.

torchref.config.normalize_device(dev=None)

Coerce a user-supplied device (or None) to a canonical device.

None resolves to get_default_device(). The pure, side-effect-free counterpart of torchref.utils.resolve_device(), which moves the objects it is given: use this for one device source, that one to reconcile several.

class torchref.config.DeviceConfig

The active device: device.current reads it, assignment sets it.

Resolved once at import cuda -> mps -> cpu by _auto_detect_device(), which gates CUDA on compute capability and _MIN_CUDA_VRAM_GB of VRAM; TORCHREF_DEVICE overrides that, bypassing the gates but raising if the backend is unavailable. The setter mirrors it – a bad value raises ValueError/RuntimeError rather than silently falling back, so callers can decide how to recover.

__init__()

property current: device

Get the current default device.

torchref.config.get_default_device()

Get the current default device.