File size: 2,322 Bytes
c65e212 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """Scalar schedules for structure-steering potentials."""
from __future__ import annotations
import math
from abc import ABC, abstractmethod
from collections.abc import Sequence
class ParameterSchedule(ABC):
"""Map normalized diffusion time ``t`` to a potential parameter."""
@abstractmethod
def compute(self, t: float) -> float:
"""Evaluate the schedule at ``t``."""
class ExponentialInterpolation(ParameterSchedule):
"""Interpolate from ``start`` to ``end`` with exponential curvature."""
def __init__(self, start: float, end: float, alpha: float) -> None:
self.start = start
self.end = end
self.alpha = alpha
def compute(self, t: float) -> float:
span = self.end - self.start
if self.alpha == 0:
return self.start + span * t
numerator = math.exp(self.alpha * t) - 1
denominator = math.exp(self.alpha) - 1
return self.start + span * numerator / denominator
class PiecewiseStepFunction(ParameterSchedule):
"""Select values separated by strict upper thresholds.
A time exactly equal to a threshold remains in the lower interval. This
boundary convention is part of the steering-input contract.
"""
def __init__(
self,
thresholds: Sequence[float],
values: Sequence[float],
) -> None:
self.thresholds = tuple(thresholds)
self.values = tuple(values)
if not self.thresholds:
raise ValueError("PiecewiseStepFunction requires at least one threshold.")
if len(self.values) != len(self.thresholds) + 1:
raise ValueError(
"PiecewiseStepFunction requires exactly one more value than threshold; "
f"received {len(self.values)} values and {len(self.thresholds)} thresholds."
)
if any(
current >= following
for current, following in zip(self.thresholds, self.thresholds[1:], strict=False)
):
raise ValueError("PiecewiseStepFunction thresholds must be strictly increasing.")
def compute(self, t: float) -> float:
interval = next(
(index for index, threshold in enumerate(self.thresholds) if t <= threshold),
len(self.thresholds),
)
return self.values[interval]
|