From d4a40d72b374a0e56cc66e3b4ec8ac263a047d7f Mon Sep 17 00:00:00 2001 From: Charlelie Laurent Date: Fri, 16 Oct 2026 13:13:54 -0700 Subject: [PATCH] Add inverse SNR and TrigFlow components Signed-off-by: Charlelie Laurent --- CHANGELOG.md | 10 + docs/api/diffusion/noise_schedulers.rst | 40 ++++ docs/api/diffusion/preconditioners.rst | 8 + .../diffusion/noise_schedulers/__init__.py | 1 + physicsnemo/diffusion/noise_schedulers/edm.py | 11 +- .../diffusion/noise_schedulers/iddpm.py | 7 +- .../noise_schedulers/linear_gaussian.py | 93 +++++++- .../noise_schedulers/rectified_flow.py | 9 +- .../diffusion/noise_schedulers/trigflow.py | 190 +++++++++++++++++ physicsnemo/diffusion/noise_schedulers/ve.py | 8 +- physicsnemo/diffusion/noise_schedulers/vp.py | 18 +- .../diffusion/preconditioners/__init__.py | 1 + .../preconditioners/preconditioners.py | 88 ++++++++ .../data/trigflow_precond_conv.mdlus | Bin 0 -> 3608 bytes .../trigflow_precond_conv_coefficients.pth | Bin 0 -> 2731 bytes .../data/trigflow_precond_conv_forward.pth | Bin 0 -> 4163 bytes .../data/trigflow_precond_conv_sigma.pth | Bin 0 -> 1845 bytes .../data/trigflow_precond_linear.mdlus | Bin 0 -> 3998 bytes .../trigflow_precond_linear_coefficients.pth | Bin 0 -> 2751 bytes .../data/trigflow_precond_linear_forward.pth | Bin 0 -> 2129 bytes .../data/trigflow_precond_linear_sigma.pth | Bin 0 -> 1859 bytes test/diffusion/test_model_io.py | 4 +- test/diffusion/test_noise_schedulers.py | 4 + test/diffusion/test_preconditioners.py | 6 +- test/diffusion/test_trigflow.py | 198 ++++++++++++++++++ 25 files changed, 683 insertions(+), 13 deletions(-) create mode 100644 physicsnemo/diffusion/noise_schedulers/trigflow.py create mode 100644 test/diffusion/data/trigflow_precond_conv.mdlus create mode 100644 test/diffusion/data/trigflow_precond_conv_coefficients.pth create mode 100644 test/diffusion/data/trigflow_precond_conv_forward.pth create mode 100644 test/diffusion/data/trigflow_precond_conv_sigma.pth create mode 100644 test/diffusion/data/trigflow_precond_linear.mdlus create mode 100644 test/diffusion/data/trigflow_precond_linear_coefficients.pth create mode 100644 test/diffusion/data/trigflow_precond_linear_forward.pth create mode 100644 test/diffusion/data/trigflow_precond_linear_sigma.pth create mode 100644 test/diffusion/test_trigflow.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b76c9207ab..bc8d798286 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 as in sCM and EDM2. Existing single-time, single-prediction calls remain unchanged. Existing DSM and flow-matching losses continue to consume single tensors. +- Adds `TrigFlowNoiseScheduler` and `TrigFlowPreconditioner` to the diffusion + noise schedulers and preconditioners in `physicsnemo.diffusion`. Pair them + with the same `sigma_data` to train and sample models along a trigonometric + diffusion path, with the backbone predicting normalized velocity. Built-in + linear Gaussian schedulers also provide differentiable `snr_inv()` methods + to convert signal-to-noise amplitude ratios back to diffusion times. Use + these inverses to match noise levels across schedules. The amplitude ratio + is `alpha(t) / sigma(t)`. The power signal-to-noise ratio is its square. + Custom schedulers can supply their own inverse. These additions are opt-in. + Existing schedulers, preconditioners, and training workflows work as before. - Adds standalone FLARE++ attention and model APIs, with input-conditioned dynamic routing, plus a ``GALE_FPP`` backend for using the same mixer inside GeoTransolver. diff --git a/docs/api/diffusion/noise_schedulers.rst b/docs/api/diffusion/noise_schedulers.rst index dda757e7fe..90f0d58675 100644 --- a/docs/api/diffusion/noise_schedulers.rst +++ b/docs/api/diffusion/noise_schedulers.rst @@ -23,6 +23,46 @@ the coupling between noise schedule and reverse process in one place and avoids the need for schedule-specific sampler code. +Amplitude ratio and inverse +--------------------------- + +The amplitude ratio :math:`\alpha(t)/\sigma(t)` measures signal strength +relative to noise strength. The method ``scheduler.snr(t)`` returns this +ratio. The usual power signal-to-noise ratio is its square. + +Use ``scheduler.snr_inv(q)`` to recover a diffusion time from an amplitude +ratio or to match noise levels between schedules. EDM, VE, VP, rectified-flow, +continuous iDDPM, and TrigFlow provide differentiable inverses. Use finite +positive ratios on the schedule's invertible domain. Inverses do not clip +to training bounds. Add ``snr_inv`` for custom schedulers that need this +conversion, and keep it differentiable for training losses such as sCM. + +.. code-block:: python + + import torch + from physicsnemo.diffusion.noise_schedulers import TrigFlowNoiseScheduler + + scheduler = TrigFlowNoiseScheduler(sigma_data=0.5) + t = torch.tensor([0.2, 0.7]) + recovered = scheduler.snr_inv(scheduler.snr(t)) + torch.testing.assert_close(recovered, t) + +TrigFlow interpolates between data and Gaussian noise with trigonometric +coefficients. Its diffusion time is an angle in radians, whereas EDM uses +the noise standard deviation. + +Pair ``TrigFlowNoiseScheduler`` with ``TrigFlowPreconditioner`` and the same +``sigma_data``. Training samples diffusion times uniformly. Inference uses +an evenly spaced grid ending at zero. The denoising weight +:math:`1/(\sigma_d\sin(t))^2` gives unit weight to the normalized-velocity +residual with this preconditioner. Keep training diffusion times above zero +for denoising losses and below :math:`\pi/2` for reverse-process sampling. + +.. autoclass:: TrigFlowNoiseScheduler + :show-inheritance: + :members: + + Role in Training ---------------- diff --git a/docs/api/diffusion/preconditioners.rst b/docs/api/diffusion/preconditioners.rst index 8a80c59f24..5f150a6c94 100644 --- a/docs/api/diffusion/preconditioners.rst +++ b/docs/api/diffusion/preconditioners.rst @@ -140,6 +140,14 @@ API Reference :members: :exclude-members: forward +:code:`TrigFlowPreconditioner` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: physicsnemo.diffusion.preconditioners.TrigFlowPreconditioner + :show-inheritance: + :members: + :exclude-members: forward + :code:`VEPreconditioner` ~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/physicsnemo/diffusion/noise_schedulers/__init__.py b/physicsnemo/diffusion/noise_schedulers/__init__.py index 21cbe392a6..5316af9fc3 100644 --- a/physicsnemo/diffusion/noise_schedulers/__init__.py +++ b/physicsnemo/diffusion/noise_schedulers/__init__.py @@ -21,6 +21,7 @@ from .linear_gaussian import LinearGaussianNoiseScheduler # noqa: F401 from .rectified_flow import RectifiedFlowNoiseScheduler # noqa: F401 from .student_t_edm import StudentTEDMNoiseScheduler # noqa: F401 +from .trigflow import TrigFlowNoiseScheduler # noqa: F401 from .ve import VENoiseScheduler # noqa: F401 from .vp import VPNoiseScheduler # noqa: F401 diff --git a/physicsnemo/diffusion/noise_schedulers/edm.py b/physicsnemo/diffusion/noise_schedulers/edm.py index 6b0174e103..5189965ea3 100644 --- a/physicsnemo/diffusion/noise_schedulers/edm.py +++ b/physicsnemo/diffusion/noise_schedulers/edm.py @@ -71,7 +71,9 @@ class EDMNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the EDM noise scheduler: + The first sequence covers the EDM training and sampling interfaces with one + scheduler. Its final call converts a clean-data predictor into the right-hand side + used by an ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler @@ -98,7 +100,8 @@ class EDMNoiseScheduler(LinearGaussianNoiseScheduler): >>> denoiser(xN, tN).shape # ODE RHS for sampling torch.Size([4, 3, 8, 8]) - Per-channel ``sigma_data`` for heterogeneous channels: + A channel-specific ``sigma_data`` assigns different data scales to each channel. + The resulting loss weight retains a separate value for each channel. >>> sigma_per_ch = torch.tensor([0.3, 0.5, 0.7]) >>> scheduler_ch = EDMNoiseScheduler(sigma_data=sigma_per_ch) @@ -143,6 +146,10 @@ def sigma_inv( r"""Identity mapping: :math:`t = \sigma`.""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = 1/q` for amplitude ratios ``q > 0``.""" + return self.sigma_inv(q.reciprocal()) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/iddpm.py b/physicsnemo/diffusion/noise_schedulers/iddpm.py index ca33329988..126a5fc589 100644 --- a/physicsnemo/diffusion/noise_schedulers/iddpm.py +++ b/physicsnemo/diffusion/noise_schedulers/iddpm.py @@ -59,7 +59,8 @@ class IDDPMNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the iDDPM noise scheduler: + The first sequence builds noisy training samples and a reverse-time sampling state + for iDDPM. It then wraps a clean-data predictor as an ODE right-hand side. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import IDDPMNoiseScheduler @@ -131,6 +132,10 @@ def sigma_inv( r"""For iDDPM, :math:`t = \sigma` (identity mapping).""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find continuous diffusion time :math:`t = 1/q` for amplitude ratios ``q > 0``.""" + return q.reciprocal() + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py b/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py index ca5d7ae3e0..65196bbd97 100644 --- a/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py +++ b/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py @@ -130,6 +130,7 @@ class LinearGaussianNoiseScheduler(ABC, NoiseScheduler): ... smin_rho = self.sigma_min**(1/self.rho) ... frac = i/(num_steps-1) ... t = (smax_rho + frac * (smin_rho - smax_rho))**self.rho + ... # Include the clean endpoint after the reverse-time grid. ... return torch.cat([t, torch.zeros(1, device=device)]) ... ... def sample_time(self, N, *, device=None, dtype=None): @@ -456,6 +457,50 @@ def snr( """ return self.alpha(t) / self.sigma(t) + def snr_inv( + self, + q: Float[Tensor, " *shape"], + ) -> Float[Tensor, " *shape"]: + r"""Find the diffusion time for a signal-to-noise amplitude ratio. + + Use this to match noise levels between diffusion schedules. The + ratio compares signal and noise amplitudes. The usual power + signal-to-noise ratio is its square. + + .. math:: + + q = \alpha(t) / \sigma(t), \qquad t = \mathrm{snr\_inv}(q). + + Parameters + ---------- + q : Tensor + Amplitude ratio, not its square. Built-in continuous inverses + support finite positive ratios without clipping to training bounds. + + Returns + ------- + Tensor + Diffusion times with the shape of ``q``. + + Notes + ----- + Override this method if your custom scheduler needs an inverse. + Keep the calculation differentiable for training losses such as sCM. + Schedulers without an inverse raise ``NotImplementedError`` here. + + Examples + -------- + For EDM, the signal coefficient is one and the noise coefficient equals time. + Inverting a signal-to-noise ratio of two gives time one half. + + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler + >>> EDMNoiseScheduler().snr_inv(torch.tensor([2.0])) + tensor([0.5000]) + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement snr_inv()." + ) + def x0_to_score( self, x0: Float[Tensor, " B *dims"], @@ -491,6 +536,10 @@ def x0_to_score( Examples -------- + A caller can convert a clean-data predictor into a score predictor when it needs + score parameterization. The conversion uses the noisy state and time to recover + the implied Gaussian score. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion @@ -562,6 +611,10 @@ def score_to_x0( Examples -------- + A caller can convert a score predictor into a clean-data predictor for a loss or + denoiser that expects that parameterization. The conversion combines the score + with the noisy state at the same time. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion @@ -620,6 +673,10 @@ def epsilon_to_score( Examples -------- + A caller can convert a noise predictor into a score predictor without changing + its backbone. The scheduler applies its time-dependent noise coefficient to the + predicted noise. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an epsilon-predictor, wrap it for manual conversion @@ -677,6 +734,9 @@ def score_to_epsilon( Examples -------- + A score predictor can also supply a noise prediction for a noise-based consumer. + The scheduler scales the score by its time-dependent noise coefficient. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion: @@ -740,6 +800,10 @@ def epsilon_to_x0( Examples -------- + A noise predictor can supply a clean-data estimate for losses that train on + clean targets. The conversion subtracts predicted noise from the noisy state + using the schedule coefficients. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an epsilon-predictor, wrap it for manual conversion @@ -801,6 +865,10 @@ def x0_to_epsilon( Examples -------- + A clean-data predictor can supply a noise estimate without retraining the + backbone. The conversion subtracts the predicted clean signal from the noisy + state. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion: @@ -867,6 +935,10 @@ def x0_to_flow( Examples -------- + A clean-data predictor can supply a velocity target for a flow-matching loss. + The conversion differentiates the schedule coefficients at the same diffusion + time. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion @@ -934,6 +1006,10 @@ def flow_to_x0( Examples -------- + A caller can convert a velocity predictor to a clean-data predictor when it + needs the original sample. The scheduler solves the linear-Gaussian velocity + relation using the noisy state. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a flow-predictor, wrap it for manual conversion: @@ -1003,6 +1079,10 @@ def score_to_flow( Examples -------- + A score predictor can supply the velocity expected by a flow-matching loss. The + scheduler first relates the score to the clean signal, then applies the schedule + derivatives. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion @@ -1073,6 +1153,9 @@ def flow_to_score( Examples -------- + A velocity predictor can supply the score needed by a denoiser. The scheduler + uses the noisy state and schedule derivatives to perform the conversion. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a flow-predictor, wrap it for manual conversion @@ -1190,7 +1273,9 @@ def get_denoiser( Examples -------- - Generate ODE RHS from a score-predictor: + A denoiser returns the reverse-process right-hand side used by a solver. This + first example starts from a score predictor and requests the probability-flow + ODE. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler @@ -1204,7 +1289,8 @@ def get_denoiser( >>> dx_dt.shape torch.Size([2, 3, 8, 8]) - Generate ODE RHS from an x0-predictor (score conversion is done internally): + A clean-data predictor can use the same ODE interface without producing a + score explicitly. The scheduler converts its output to the required score. >>> x0_pred = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2) # Toy x0-predictor >>> denoiser = scheduler.get_denoiser( @@ -1213,7 +1299,8 @@ def get_denoiser( >>> dx_dt.shape torch.Size([2, 3, 8, 8]) - Generate ODE RHS from an epsilon-predictor: + A noise predictor can also drive the ODE through the same factory. + The scheduler converts its output before evaluating the right-hand side. >>> eps_pred = lambda x, t: x * 0.1 # Toy epsilon-predictor >>> denoiser = scheduler.get_denoiser( diff --git a/physicsnemo/diffusion/noise_schedulers/rectified_flow.py b/physicsnemo/diffusion/noise_schedulers/rectified_flow.py index 9ec1935f1d..f21eba4539 100644 --- a/physicsnemo/diffusion/noise_schedulers/rectified_flow.py +++ b/physicsnemo/diffusion/noise_schedulers/rectified_flow.py @@ -86,8 +86,9 @@ class RectifiedFlowNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Construct noisy states for training, then initialize sampling with the same - flow parameterization: + Training interpolates clean data toward noise, while sampling starts near the noisy + endpoint. The flow predictor then supplies the ODE right-hand side under the same + time convention. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import ( @@ -154,6 +155,10 @@ def sigma_inv( r"""Identity mapping: :math:`t = \sigma`.""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = 1/(1+q)` for amplitude ratios ``q > 0``.""" + return 1 / (1 + q) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/trigflow.py b/physicsnemo/diffusion/noise_schedulers/trigflow.py new file mode 100644 index 0000000000..69e96759c3 --- /dev/null +++ b/physicsnemo/diffusion/noise_schedulers/trigflow.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Trigonometric interpolation between data and Gaussian noise.""" + +import math + +import torch +from jaxtyping import Float +from torch import Tensor + +from .linear_gaussian import LinearGaussianNoiseScheduler + + +class TrigFlowNoiseScheduler(LinearGaussianNoiseScheduler): + r"""Interpolate between clean data and Gaussian noise along a trigonometric path. + + Use this scheduler for TrigFlow denoising or consistency training, paired + with ``TrigFlowPreconditioner`` and the same ``sigma_data``. Diffusion time + is an angle in radians, whereas EDM uses the noise standard deviation. + Training draws diffusion times uniformly. Inference uses evenly spaced + diffusion times ending at zero. + + .. math:: + + x_t = \cos(t)x_0 + \sigma_d \sin(t)\epsilon, + \qquad \epsilon \sim \mathcal{N}(0,I). + + Parameters + ---------- + sigma_data : float, optional, default=0.5 + Data scale and noise standard deviation at :math:`t=\pi/2`. + Must be finite and positive. + t_min : float, optional, default=0.001 + Lower training diffusion time in radians, with + :math:`0 \leq t_{\min} < t_{\max}`. Use a positive value for + denoising loss weights, which diverge at zero. + t_max : float, optional, default=1.55 + Upper training diffusion time and first inference diffusion time, + in radians. Must be below ``pi / 2``. The default stays below that + endpoint in reduced precision. + + Examples + -------- + Use this scheduler when the model treats diffusion time as an angle and predicts + TrigFlow velocity. The example checks noisy-state shape, reverse-time grid size, and + the inverse signal-to-noise mapping. + + >>> scheduler = TrigFlowNoiseScheduler() + >>> t = scheduler.sample_time(2) + >>> scheduler.add_noise(torch.zeros(2, 3), t).shape + torch.Size([2, 3]) + >>> scheduler.timesteps(4).shape + torch.Size([5]) + >>> # Use interior angles for the inverse SNR check. + >>> t = torch.tensor([0.2, 0.7]) + >>> torch.allclose(scheduler.snr_inv(scheduler.snr(t)), t) + True + + Note + ---- + Coefficient functions accept angular times without clipping. The training + bounds do not limit coefficient evaluation. At zero, score conversions + and denoising weights are singular. At ``pi / 2``, conversions that divide + by the signal coefficient are singular. + """ + + def __init__( + self, sigma_data: float = 0.5, t_min: float = 0.001, t_max: float = 1.55 + ) -> None: + if not math.isfinite(sigma_data) or sigma_data <= 0: + raise ValueError(f"Expected finite sigma_data > 0, got {sigma_data}.") + if not 0 <= t_min < t_max < math.pi / 2: + raise ValueError( + f"Expected 0 <= t_min < t_max < pi/2, got {t_min=}, {t_max=}." + ) + self.sigma_data = sigma_data + self.t_min = t_min + self.t_max = t_max + + def alpha(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Signal coefficient: :math:`\alpha(t)=\cos(t)`.""" + return t.cos() + + def alpha_dot(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Signal derivative: :math:`\dot\alpha(t)=-\sin(t)`.""" + return -t.sin() + + def sigma(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Noise coefficient: :math:`\sigma(t)=\sigma_d\sin(t)`.""" + return self.sigma_data * t.sin() + + def sigma_dot(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Noise derivative: :math:`\dot\sigma(t)=\sigma_d\cos(t)`.""" + return self.sigma_data * t.cos() + + def sigma_inv(self, sigma: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`\arcsin(\sigma/\sigma_d)` for noise scales in ``[0, sigma_data]``.""" + return (sigma / self.sigma_data).asin() + + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`\arctan(1/(\sigma_d q))` for nonnegative amplitude ratios.""" + return torch.atan2(torch.ones_like(q), self.sigma_data * q) + + def sample_time( + self, + N: int, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> Float[Tensor, " N"]: + r"""Draw training diffusion times uniformly between ``t_min`` and ``t_max``. + + Parameters + ---------- + N : int + Number of training examples. + device : torch.device, optional + Device for the sampled times. + dtype : torch.dtype, optional + Floating-point type for the sampled times. + + Returns + ------- + Tensor + Training diffusion times in radians with shape :math:`(N,)`. + """ + return self.t_min + (self.t_max - self.t_min) * torch.rand( + N, device=device, dtype=dtype + ) + + def timesteps( + self, + num_steps: int, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> Float[Tensor, " N+1"]: + r"""Return evenly spaced inference diffusion times from ``t_max`` to zero. + + Parameters + ---------- + num_steps : int + Number of solver steps, at least one. + device : torch.device, optional + Device for the inference grid. + dtype : torch.dtype, optional + Floating-point type for the grid. + + Returns + ------- + Tensor + Decreasing diffusion times in radians with shape + :math:`(\mathrm{num\_steps}+1,)`. + """ + if num_steps < 1: + raise ValueError(f"Expected num_steps >= 1, got {num_steps}.") + return torch.linspace(self.t_max, 0, num_steps + 1, device=device, dtype=dtype) + + def loss_weight(self, t: Float[Tensor, " N"]) -> Float[Tensor, " N"]: + r"""Weight denoising residuals by :math:`1/(\sigma_d\sin(t))^2`. + + Use with ``TrigFlowPreconditioner`` to give unit weight to the + normalized-velocity residual. Use positive diffusion times to avoid + the singularity at zero. + + Parameters + ---------- + t : Tensor + Training diffusion times in radians with shape :math:`(N,)`. + + Returns + ------- + Tensor + Per-example denoising weights. + """ + return self.sigma(t).square().reciprocal() diff --git a/physicsnemo/diffusion/noise_schedulers/ve.py b/physicsnemo/diffusion/noise_schedulers/ve.py index 608bbdc1de..4351b8faea 100644 --- a/physicsnemo/diffusion/noise_schedulers/ve.py +++ b/physicsnemo/diffusion/noise_schedulers/ve.py @@ -55,7 +55,9 @@ class VENoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the VE noise scheduler: + The first sequence uses a variance-exploding schedule for both training and + sampling. The final predictor conversion supplies the right-hand side expected by an + ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import VENoiseScheduler @@ -105,6 +107,10 @@ def sigma_inv( r"""Inverse VE mapping: :math:`t = \sigma^2`.""" return sigma**2 + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = q^{-2}` for amplitude ratios ``q > 0``.""" + return self.sigma_inv(q.reciprocal()) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/vp.py b/physicsnemo/diffusion/noise_schedulers/vp.py index ee6582822d..010f4c5cfd 100644 --- a/physicsnemo/diffusion/noise_schedulers/vp.py +++ b/physicsnemo/diffusion/noise_schedulers/vp.py @@ -66,7 +66,9 @@ class VPNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the VP noise scheduler: + The first sequence uses a variance-preserving schedule for both training and + sampling. The final predictor conversion supplies the right-hand side expected by an + ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import VPNoiseScheduler @@ -158,6 +160,20 @@ def sigma_inv( discriminant = self.beta_min**2 + 2 * self.beta_d * a return 2 * a / (self.beta_min + torch.sqrt(discriminant)) + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time for finite amplitude ratios ``q > 0``. + + .. math:: + + A = \log(1+q^{-2}), \qquad + t = \frac{2A}{\beta_{\min}+\sqrt{\beta_{\min}^2+2\beta_d A}}. + + The rationalized root avoids cancellation and supports ``beta_d=0`` + when ``beta_min > 0``. See the base ``snr_inv`` contract. + """ + a = torch.logaddexp(torch.zeros_like(q), -2 * q.log()) + return 2 * a / (self.beta_min + (self.beta_min**2 + 2 * self.beta_d * a).sqrt()) + def timesteps( self, num_steps: int, diff --git a/physicsnemo/diffusion/preconditioners/__init__.py b/physicsnemo/diffusion/preconditioners/__init__.py index 848114c5f8..8dcdabe631 100644 --- a/physicsnemo/diffusion/preconditioners/__init__.py +++ b/physicsnemo/diffusion/preconditioners/__init__.py @@ -28,6 +28,7 @@ BaseAffinePreconditioner, EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) diff --git a/physicsnemo/diffusion/preconditioners/preconditioners.py b/physicsnemo/diffusion/preconditioners/preconditioners.py index eeac45dd74..3a5392ea02 100644 --- a/physicsnemo/diffusion/preconditioners/preconditioners.py +++ b/physicsnemo/diffusion/preconditioners/preconditioners.py @@ -580,6 +580,94 @@ def forward( return c_skip * x + c_out * F_x +class TrigFlowPreconditioner(BaseAffinePreconditioner): + r"""Convert normalized-velocity predictions into clean-data predictions. + + Use with ``TrigFlowNoiseScheduler`` and the same ``sigma_data`` for + TrigFlow denoising or consistency training. The backbone predicts velocity + divided by ``sigma_data``, rather than clean data or noise. Diffusion times + are angles in radians. A backbone can return a prediction paired with + ``DiffusionExtraOutputs``; its learned loss weight passes through unchanged. + + .. math:: + + D_\theta(x,t) = \cos(t)x - \sigma_d\sin(t)F_\theta(x/\sigma_d,t). + + Parameters + ---------- + model : physicsnemo.Module + Backbone predicting velocity divided by ``sigma_data``. + sigma_data : float, optional, default=0.5 + Data scale, matching the noise scheduler. Must be finite and positive. + + Forward + ------- + x : Tensor + Noisy states with shape :math:`(B, \ldots)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time in radians. Pass ``(t, extra_times)`` to include a + destination diffusion time ``r``. + condition : Tensor or TensorDict, optional + Conditioning passed to the backbone. + **model_kwargs : Any + Model-specific keyword arguments. + + Outputs + ------- + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + Clean-data prediction, optionally paired with its unchanged loss weight. + + Examples + -------- + The backbone in this example predicts zero normalized velocity, so the wrapper + returns the clean-data skip term. At zero angular time, that term equals the input. + + >>> class Velocity(Module): + ... def forward(self, x, t, condition=None): + ... return torch.zeros_like(x) + >>> model = TrigFlowPreconditioner(Velocity()) + >>> model(torch.ones(2, 3), torch.zeros(2)) + tensor([[1., 1., 1.], + [1., 1., 1.]]) + """ + + sigma_data: Float[torch.Tensor, ""] + + def __init__(self, model: Module, sigma_data: float = 0.5) -> None: + super().__init__(model) + if not math.isfinite(sigma_data) or sigma_data <= 0: + raise ValueError(f"Expected finite sigma_data > 0, got {sigma_data}.") + self.register_buffer("sigma_data", torch.tensor(float(sigma_data))) + + def compute_coefficients( + self, t: Float[torch.Tensor, " B *dims"], / + ) -> tuple[ + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + ]: + r"""Return input, time, output, and skip coefficients for diffusion time. + + Parameters + ---------- + t : Tensor + Diffusion times in radians with shape :math:`(B, 1, \ldots)` + for broadcasting over the noisy states. + + Returns + ------- + tuple of Tensor + Coefficients ``(1 / sigma_data, t, -sigma_data * sin(t), cos(t))``. + """ + return ( + torch.ones_like(t) / self.sigma_data, + t, + -self.sigma_data * t.sin(), + t.cos(), + ) + + class VPPreconditioner(BaseAffinePreconditioner): r""" Variance Preserving (VP) preconditioner. diff --git a/test/diffusion/data/trigflow_precond_conv.mdlus b/test/diffusion/data/trigflow_precond_conv.mdlus new file mode 100644 index 0000000000000000000000000000000000000000..d5ba176859709653ac8d17e8b8350739fc5b4034 GIT binary patch literal 3608 zcmbtX4RBLc7Jg~^o3`|yLcxNiEXA#_`RfnZyqgIU&^OT9MH*xAktVOnqxtE}dnvNr zteuf+m9;twyQ2c4I}Yy1kDw#!ns;j*6c^mxqSGk?G6?SCZy^gY4?`S}ijOyKMBPIGq8nmvxF~T7v|?jqtVR<1r%idW?CiK(k5*46=8mXCpw2&6077Z!~DYNd`M^~o9T?q z1E2qk-i<K8MJ}En6-5iru%{b(s0k*r^BcUSOG;xs0fvWR7nw zx+Bb=804J-}a@&pAo$$uX ziK3UqKJ=N+%}z)5u}tm541!AdfH%8FYqc!>84ef>uL5lSlzP*4bdU)w(^)smt?R6o z#-Cx`Ztt7(1bnpjDL6v+O7|`NR4Q4157@j&2Ykn#mxh>(!3Ncf)TP&+hu5^PQg6Mx zgW4DvH$Bz46AWA&0QS1DG_+tZTv2=irsNlZ`xek(%c|Mn-rV=$`|s_OR@Zz3HsqK= z<+&1gbwDTm@$7BX;K%`>$~_9Un}fjkpH_IJ=qBiLt4z0;uSi|%e@hL|J4=n!|HJfy zrV=!sIZM6y@Ni`N={rGF(^jdpu+>CEixZUeH-qabre3=R1H7n88BY*FgQH7 zj>&SQUv#M?P=WY4(NI&w|=@X3DI`uYH< z&dHa~EME=oy%k0t`m6?wJi8Vy{D51^Jk$!-3?Gw9bNj$3bBNmW&mPI&R%F(l{0=^` zkB*#rrw=xq7U65QEK}o=T5$N~;~=$cA3VF^6)JaT8rW-IG+6SB6HxmGCoNn1IA#6h zH`G}74%6(sz0^IsFHn`2zMwAkuYf2fr&>FU=jDGg!A~TjY`2?8wvZZa}_w zkQzJR345zMBERJ3!T)-{kuF~u2f?zxO4|mnns$|L1GVj5aQ?{y;LFN(_)*V@G%ss4 z96Fj0PMYU}w(K#gV*kIvj-q~exou~pv;AAs;k;Qecgb_``;y^EGrbua3U^7m{XNv5 z$2Lm`?H5g%-e;&wM@~jYo_PRlX}u=hGHL;rSC)gXRz59#zS1n!ogR;n18-2vKG-GA zDS8xU{-K-Ndgvg`9=KrI?S3qJ2JR^Tx~3mp0a&Ix1C8=Cpto8YZukr=((Ix7f7c5h zzcY3YN7L7w%Ry%l%hW#PC0&Jy%k_rtKZ6FJ{^6bOyMOxrD<|}cm6OrSM__{-x`BV~ zz}1`KTviMDe?~&vpsh1%bq2k5iNUDX>U6sLdYw*R-_TH7r>`{_v`dWodZST~y=-cF z=}%QCE?Amr6?Qn$c_bGnT&?9p*@r06$pxrX2EB}G(ga#j&Xgp{MWjqdau;$nL2_N? zQt~7hf&v-O@g$x^0aC(~Z*tsWOpdjv&zW!|Ry&l~D6&ex%$=w8n;4tO#H5h zCL2NAA8&LK4mX$t#LX!hvCbt8*KUN#BCvi;7a!uuHkzpX=x?_pWh|K65p)aW1|i_@ zsbCnt!^bjA#nSl|Yl55`HLDw&gRC>)cX9Z5vB8S!`A7?OXd!Q0ugkMZ;GBY=^#w>5 z*U=He4IbGgLkY_F7!9@@HP%xR?fI}-hDjz|6*ol2{zt|qyKtnfBaPvxL(^cE&X=!i zS-{_oXj$*|OA%Q?Oj-keWfTj!&T*&5;rFv%0eNExhr{wvn_pS?Kp2G>OD>X$_BbBy zB(Ep>Qyc<4X&|)@EjHzIc|!uz!J%NsGY>0hNgWyed%9qe(_xXHs+3~&WKt-ajAl`>8C|zElyc=Uk-9rBuW~^slrTElnNXYWKB`Id76Svfk3VWkJs?Xd`WlSM^0t7^q$Z`?BzC&a0!eS+N59Ei48O7%ab)HE3d; zGgw0i4rk%{YA};-@jYF9!>r>~Bt6qGPH5nB5+}i;qR!!C79j@M1^BB8i<92iz&IFoVqoSW5)=|V zsieQ5?rx??B0`ctJUqFy-9ug>>5V;%fFUOyNLEgSQpth+49MSNtho~>e=}T9D|SKp z|D5zVb`nk{JL@}n^{JhG`|K*~{^DgjXJ@_<16>$bW%9Co!iG@_My`F(^gh|BWOd=kNMVS~wPsW+aP%M$4 znGj7gF*+K-0}>HEOE-H5mjgH^NzbHR#BgorGS{jx^1#v#u#efn=GsJCpaE>a0oYdf z?0j?SiCXw;4*7fPon69QS;j3AZFp~VFojx0vy;u0V#Gpz?u=~IDLd9&0IXO)J}5Xa ztYh)kVvP}VUR$vn2f)&HtT|n+SU=v~ow{Z0SaVK}TGn>M8FnZ0?ObyfB`n-Fy!AM+ zTR*=uX|H!Txcc1}YhY^UZufhQ9LI3`K=wJ@@c+nR4qV$AV4MM(^`{h|02yzL{Rd_t BrUC!} literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_conv_forward.pth b/test/diffusion/data/trigflow_precond_conv_forward.pth new file mode 100644 index 0000000000000000000000000000000000000000..b38b7888001a44ffc59d68b5680a878cffc1a2b5 GIT binary patch literal 4163 zcmbVPc|4VA8$Nc95K<9}B(lu0gp%_-$2RIvG_*%LPIV+Y9BHM^GDf4M(n(4oBc-yH z+)vH$kJ$vinPBi zW7-s&Ae^^!Nk~{!z{+qgCoD85Kzvyhup}%zDlk0gyP!ZppykTtA@0r&w0ZyM+9p(43W#e0eV{E#NEouxTc2S+)#Yku5La zE1Iykvl(o8wv>Rcl)&MHg@kZ90$x~X1bs$$5I39~Cb}x;!E+`pi9FV7s>@JBGy^z9Lt`TsmRvjV^7Ng*m8gS;6I%J!a zB7@xD(CwaSu()C@IaZWH_8k{umn&9y`rebsDK8$*?P?N*MU}$r2PwpFlN>dqXe86F z_5$YnsImT~EJsZ|)+BOD(!(VI_Bc(c4Ni8nK)d!@()H7;%5KB6K)1+dzIvzv{WBC{ zaZ5Bh9#n?4n#v#{Di$`)qEr53qG6{-KPh)UflKDMqs?RAliJrgq#@oD&)?xjJ*>^b z2+6WC@{_?^V1~zx-HF%VI}2lNONiOfG^pC~01I}#6gf}W0kxs!V5Khw<(0RfX-ghN zj?jg)R1Ub5JV(u=je+}%GWEE}5AxP5M$8Q=tfy|x_^Ng={2KlOx)rn0>hpcjzaWjc zj&mbZ7*DYM3?EkHuYKXAoy>elV?RXC6NbZNcMT-SJxVf|>k0i; z6r3+h0lMJI`D8Fe)Whw2$4g7exU}gbUs(a8XUNU(L8{rUVBHCzp9hGh#4%^>7 z0TKI0n1zRRh^vW$WmRCOu6d>dJNy{L>ID`(ih_JpVF4pfbj;>mf&F1i6rxgcE{N zaCTxlxv91TEY`gwGSf?8Vy6}I_nd=Evl4;r6H0Q0iV)wO0yj3!g&Ku;;?gq>4=eZ$ zN_vH`R_F)s%hHf}?|le3eh)7A=&*8|j-$rM74XDo9HlBZj7qq8wL++BM@5;;N8Jz7 zz%soEYSKTzxZBCF;#d^$YumwBUWv6Th0dy!(g&xkMVJ|XQB=wBhts|hXzzw+*g)kX zj<-yL<7V`;?-}Mf&p6OGnGczdgdk{pO{!xiP%qxf;e)wj zsLTz~DE&b_7zHF?Tk~7!nM^5;yLtqBT(~02*;Ylg1o`->+kNDq*#XkqHsQo-p_BRh z_vHQr8)nxf4^rr%NIV?tNP@>6AX`h}kg5WB+-pX(Uo}u}>O7{h*+J&9{kKrQjqM8U4PT>|~~5kIk;Qx^yFqcc>@R z8@?s(_XG%_!1)DxOpGT%@Y00oDe(Xs7ae zk=KG&WZq_qe|6G=^A&H%v|<_9Z&*z{V=p4nh(vgotb#*&RdKc7XEk#K8_VM~)g-ywU{R&g5b>rH}#WGpkSH5YZ*nj)F@Cgge4nAnV^!y+)DIPK=}?nylJ&4xxY zra&Kt8WiE(wOdh*nlUzLPlThpyT}Rocyy((2-GyHNJ{i%n0I#%Ih*+mZPm$x1E&f} zT~Qr$#orO#=Q8lJ$l=f|6vEZL2T{$ok#LgggXH#YWO7YDv|RT=njS0RoVFoV+By?P z=pKROOUnReD8P<)aWK*(79JW=DCEQexDa{>>jv5rj!z$SWgUft*_n9eiiae~`w+Ss zvIAn=ETQ!=4?FqJAkP}ERalKUgOcK=!L?0}L@Fi%3K{{WPuvdq%NxiP|F+8Bp-E(k zr$2ION`&01ZlsC3k#62Jq7$M22s&mAF`3#=f0sxAZ@77)Z^ zKtIL9_Aa1f0sHWDb#>_I+7H7`)vOmG6 zcCJn|bzu|vA=4OQq||}$$^fDD8TN$G1!5C>zs2P(CuBL^FFZ-mFKa>-fQGuP87Jo(MI1TA(Y zw1(?Md-e@*T-OM0eRh<-?I>#P2|Y4*fh*q3i-hYl^RcO`Dpt*zhnw^Cp>}o~q_)l^ ziYjkN%_Xvn--rY#n?jQ^-KZ?HyZa7b@fY$~} zoIn6C$lZD5R|zp3Y;9~OO=Q?JtgY>At*mYA?ClwLlkDs!Svy!+iys*dR(3YFHWHRM zsU_Mb;x;AbJlOjOHJ+ z{w3t+BVs!MvVKs=|319lYwNLM# SyrrDPU5U>q&?Nbva{mp|C~MXL literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_conv_sigma.pth b/test/diffusion/data/trigflow_precond_conv_sigma.pth new file mode 100644 index 0000000000000000000000000000000000000000..777d486de2633f69ce87fa46adfaf44b1f76c5ef GIT binary patch literal 1845 zcmbuA&x_MQ6vwCApGE~i5LDQO?XFg<*)&blU12R231SY`ics_rv&pmu)25S5ySg9< zOAj9P5AdRR^We#U!RxxPCvRT#YwbHEmSd}CU>bS;J zUK>!$_BhSOTLa5BHHv$l<7z{ZBMXH>D+?RgO1*i{)k?560!g=$ z>}vD}B%`J%$=+9yni`Runi!E2uMNYoZER4>>3Q5e*TgP1Z(9cC2zRAPPEvBJNltrw zL(S%-*3O%cObW6ww`)aoE3#{R+i|-()rQ#hSW%_E`5(<=T`vLJm*3xa;b|e~qwy#o zjW%tlrrTDh*Tp@m)R*GuWC;Lj+vL{*D-JmyU=NP>BUpgkhhp%UApWoi8CHFJ6-F&VFBnQLx-7|~=az05O zEk%;d9C9%VrOZLay?335@%TDx&Bv+L}9vha$dKy1T4av_(BuYkgE@wJp@0NgyO~anG)E&YXKQ`M>}B z?mz#`{F5TLv-hG%%i5K+bi(zp>-$_NigYJDR#=ovYXrq?3bw){WXPuRDGy4i%TcLI z>F``7s}vL$XcA>(s0^!CF(N0YNTZ>1Sap$B@4Qf#OY7*|x#}F&nW3hnGQ6;KXY5$7 z&M#Cdh#O`^4qJrX3KS)*TBCPX=;-VcwI)}=(pr6yPEi`+%-Ds6r8^Q7wp8ZRjJ*ls zoTw>MvTvCNof(I89<~uOBh&5hGpD$Oag_5|#%Yn9nlfdRim;Dj8Rt}HbUGg!n7&v` zvw|{Oov&gU7Zv7F*{e)`M2xG-F^R$s#$X3>zMQ|Fal;SZdWLUu=FT!6sf=g(7%OMl zYNej>vT_x{c&i-ce3lt&^i;|C$a$5FugS|emhnqv{7cH2@#&+mCq*I6Dvg1eV3ZsX zmX1d$Ca|2DXbeIrt<&S!GLz&yeY(RSHxpFO1gnf=#Ebdkn2>TN)Hq(E&laS}9i4t^ zY+JJjM~d{}XWNIGV4iJ3UXiX)$tnhz6J_IUCfyAxRIEj^x-v(p@3Nfne}jCqAcz;1Vlz2g@nD^Zsjh#kSWq%#V)mG7bVJO zzXp3oMy3B2R8>)ny5tqyve4_Ccfvbx-5VlUczz$(s&H(H3L3$8hs%QHusUfReAl%dwSRGlOAfyem%B(|$W3q5 zUnAmv*E0*WbRB^_*K=@}R0pGePD5urC!pMXo@A=Dk1Ni4A9M!xfUcO^lIPALaQ4?d z;KYVb!>(&@!i5XBaXubt62Vsv@a;5DX|w%Fw8!@x+PiQ%y3QzIulEMnHeL+)*9thN zZD+Vl@6m8<>RGO?^cEMg0>F~SR@jh`A-&Rn8U;V!jat9I1N(%oNb;FGlB_L&GsA1( zcE1C#Gs=LflR~7evmbNG7d{22HeW*2Nf*?0u$N2gdcj@Tx)G)Qr3OxS;d5UnB*S-S zAjA3_;jnAx5)|5@j#I2pgNr-Qb3U$>@Uo&A9Q<|_r_J<~impCFTMo$$-KQ#1>@^l0 zOCKefeI^>VeflM|_dA4omLCSLV;x|F^o{aeI~>kONsM;kuQFNK0Hn!v+b1*l?ro?)6g z5dEUL&pqqwg*v~BTwP1Qq~50vMrUi_t*uAkcOlv6Z>zevz){Jl^&B5wm5zs*qaT8x z#vkEs&yDC=<{m>q_5(?qn=f*mxflK5-DyaXuS1a@^_-}275Mz&I<6_>u|#}nCHTa> z4#u?9f{;gT;ONaH?%3KN4FCM=T;O>y0j&p5Q1T@Quy;We@(?zdJTe_;~>7Hp7ZytLQWr+ zf^E%B=tAiR!=m?^xj8-U=wj`kx$Rk1F#MB|;E=ilo*F$Lmbq7<4>rf52eOmUY2s!u zwP`9$-SR&2lYE43czulS#k-+hzjcDTw=MzbcouA}4+kseX2HPzt(;%muVH59M~1!& z)u^%g47e8I3(r+efdQ%=Xxgv*;R+8IB)we+gATS!n)a4LO<)>qEvrZJeXC(%d_Q+8 zZa*s6#KOb{`{A9Esc6SRKKwSR7QBDV$?&gx8&E>}cZTMUMlL$A10E621uYv|LDDQQ z?tWY_oa3;C%bNKoG^fCwYd1KdX?+f;x;hiu^K;On6HB1S)>?R=pa(TfHlTl;_y&#a z-zhodR)qqV9*5CuZWtysc*37%&EQrBt~E@$w+?O(Rl&8TPdJ~u%h2i=FK}VY6WH6+ z4#?iKz_o3mapS?8&#U$xd+WU?B2h+Bp$U@=vKdLJueqO7_@3Z??(l|A&ubC5=)rj| zpT~92YBcO$(*$)Fec|nxXWTsIbM8v_AJO-}T!x_{dnw_c_u=Ueqk zQY~3G+d9w|ZCL>LSV14)nqUxZQ&-qJu{^~6teh;uQq4|YsvKr3Pn)d z#Gu7Gbw0i!XiPgIxnw;RN8`K zE0o^$x5*7*Hj=FhG8W&2R-v$Tm^45nF#V6ohi*aGS|o@nHwt2@V@18sm4|V#~y&dOA2%eNY#W&n6Dv8)+t^S=W3E%UzIDCiD_@?Chmjz<>CBJ>!a zP#IY*pjqQj3U=gOcuWL#<~@X6twAr>fne8l4sacu>nZecF7oJJ5d>A1MFB~wqEpwj zqNs_*%aVX7NS)8Y0|*|>!9xv2fmc*X7vbRw_E7Ou42>u}QfaVzn(SVK?O&-7L?41j zbFe?(*_fnSLvX+Xh`~VvaI6Uk7=R%Ia2&xCIXIm6HURp(RE0qcfrcjy!bp=aY7oW@ z!YKq#=ir%IILmeU-Yu`cs^ZllJ==2s;~?y8|HCC!TPPrYQPgz2hMDw*Ho_oE41L`E z`Izh$>8%YV!H`2kTam#RWT~ty;2_PU2U>Vy;=};E1b@4+IO(kogu$3I19SI~sFFa+ zq=PMVxAR1n5RygW;mM_q0NG&DTYDG>qs}~#telEu(!+ZjP<}A1xf3UUvV5QiyC8ip zCxdM`iDuG$yE*ynzQ}I&1lZ8>ZT967V2S?JhX48h){nK&_Xd>u*)Lb>?;k#^=SmZG zq945aHU7W#=@$CGZn7-9I?H}oImfQ6cN_Ef2}9lc>ra3IXY1e0OaCq#KvQJK-`>vL#px>EdxDkB1s7W{j?-xjA<)vaK-MiC3G)||OB%Mk|m;}wxvFRi|9gQ;aL_88> zm{bb?W)g{5g1CHL`5q17m?S;RxrpKB$z`r-ljMb^wRMi!(dOzDw_r)VaD=v%NV|f$ z{3I<3FB~iEsCjk~bA1V0MEv-KbTo}-U9+>z6=U4OUUkYg>XjXAE(KQb2);5nGTi{< zW5yaU<{Y|k^9TEQdu{CDn{wu8+%I$>GkHK*XMN?<3O^J>OI_T$6Ik>%Yp yKFfK%+g5{~e`!f=iEKaLk}Dc+pUByV8~%=5@W7p|0mi+cv+b`v3~(N60-mR literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_linear_forward.pth b/test/diffusion/data/trigflow_precond_linear_forward.pth new file mode 100644 index 0000000000000000000000000000000000000000..d1a46ac1ab65fee1f8d8e4c8afce0f14abb4a0a8 GIT binary patch literal 2129 zcmWIWW@cev;NW1u02&Nd3?)UG>1jFn_FwkK$VHEZzcbg&f`j-i&R9oE^!@`8heM$t9WjdBt4*MJcI8sVOd*$t7HcT$(|J z+$DuPL4~}pdk;ik`EO)%W_fB~4BSQ1~BT2u_Q*v-kD%;1YM z1UYQ``#a4b|AH`XhZ!-*6LuIhSmP6OGSl;NQ}asPoTTx&&CCSk0Q>U7mrOv*K^V6K zG#S(hIv}aCBsCut3vN!z`0Rxxww%n8lAKhK8(M82OR@p&2VvZ9P-9Rd=mt=#);Dl- zQX;|uwm0e=gc}4M*K^7`m}u_Xr=;@NJ|b$xfsG4A>>o$2winkqxc{#5e0!CpPxfb- z7TTAmraOrLdb!$D-*TDzLEZTl_KZS1)|Z`qe2`gi})UzK(YFFp4x%St|=^}y=D^I0Dcl--YY zC`dYR;Qb`k1NY~dI@Cp~A4qeaaiD``>VYZxoAxz`so2JRcX0?|JYc`6VZ{M?!FdOG zR9y}@ED^MozIt>26r1DrYdQ8GP*XKB81nG^piV^4gOh+U7sv${=cR`d zPJd<$QUqO~2af#W)S^sa4yZ~56$|m1DQ-?;I2>ePVPb4zX=-3@U}R)wVrXP+W^Qg^ zYHng=WME-wW@=_+X=ZF{Y+?X%rPBeCb`40+L4s#FoS8rXm*% z`Y1+P0ON{~sc5AXy5Y!~-U!9;i^Ld?QPQBBkDQ@3QOpkl1|1>C!wV90W06y<8j7*k zh%gpXHlUk{oc`2N%nSh5u7sQlP4(!8A}1F!6hki&VJP;59N^8yrUNxgj#(FODyT*U l0T^8lPMi!36|4|(pgyqFd#D0PO~VQ*Yk_*$89?eGY5{BgQU(A3 literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_linear_sigma.pth b/test/diffusion/data/trigflow_precond_linear_sigma.pth new file mode 100644 index 0000000000000000000000000000000000000000..76f4b47c540905b17c49eb98282cae091f29f0e0 GIT binary patch literal 1859 zcmbW2O^ee&7{{mYH7YC%f(2I)+TAX=ZQ7=3*B1goa>>D35sDsSHkqx1Y0^oiTU`+Z z=|Mc{2k_`u5d0SF!IKx+lc)U(&ZOO%WKFZ|0}mn1%>Vi2>HnGZ)JztDViBB)O>hqo z7Y+9Y z7A@6u3k~fNw22~btfmtM1|UuGpENDbG9F+HY$pKhy(Z9M|AHb{R@XhXn}rCmAY+)_ zKQ85IeU|YIWl&9EV6Lw7id;Sy!1weW>13@WRo^_5K2L$fNb1MyAJ-@a%XpFs*hnB* zI;0ej(&AXfqc6iJ$^93=jf9Fd|K&aCBI~#2cr`;9#cC*ph?Y6~6qbvG0tzTR3RcTLH_S6&uUYDpuuLUF!xD5@D8i=B|3 zNGD~rW@JjE(m})XA?c%J=`&LnRS#NuE2Q41O6=Rf+* T&J!B*vo+`673eOue!liEHA&0R literal 0 HcmV?d00001 diff --git a/test/diffusion/test_model_io.py b/test/diffusion/test_model_io.py index 4e4681b124..b1afdf679a 100644 --- a/test/diffusion/test_model_io.py +++ b/test/diffusion/test_model_io.py @@ -24,6 +24,7 @@ from physicsnemo.diffusion.preconditioners import ( EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) @@ -36,8 +37,9 @@ (VEPreconditioner, {}), (VPPreconditioner, {}), (IDDPMPreconditioner, {"M": 32}), + (TrigFlowPreconditioner, {"sigma_data": 1.4}), ) -IDS = ("edm", "ve", "vp", "iddpm") +IDS = ("edm", "ve", "vp", "iddpm", "trigflow") TIME_FIELDS = (None, (), ("r",), ("r1",), ("r2",), ("r1", "r2")) TIME_IDS = ("tensor", "empty", "r", "r1", "r2", "r1_r2") diff --git a/test/diffusion/test_noise_schedulers.py b/test/diffusion/test_noise_schedulers.py index 13fb98c917..20dbf741f7 100644 --- a/test/diffusion/test_noise_schedulers.py +++ b/test/diffusion/test_noise_schedulers.py @@ -289,6 +289,10 @@ def test_concrete_snr(self, device): t = torch.tensor([0.5, 2.0], device=device) assert torch.allclose(s.snr(t), 1 / t, atol=1e-7) + def test_optional_snr_inverse(self): + with pytest.raises(NotImplementedError, match="snr_inv"): + _MinimalScheduler().snr_inv(torch.ones(2)) + def test_concrete_add_noise(self, device): s = _MinimalScheduler() x0 = make_input((2, 3, 8, 8), seed=3, device=device) diff --git a/test/diffusion/test_preconditioners.py b/test/diffusion/test_preconditioners.py index beaccfda5d..8bee845b29 100644 --- a/test/diffusion/test_preconditioners.py +++ b/test/diffusion/test_preconditioners.py @@ -27,6 +27,7 @@ BaseAffinePreconditioner, EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) @@ -113,6 +114,7 @@ def forward( # Preconditioner configurations for parameterized tests PRECOND_CONFIGS = [ + (TrigFlowPreconditioner, {"sigma_data": 1.4}, "trigflow_precond"), ( VPPreconditioner, {"beta_d": 19.9, "beta_min": 0.1, "M": 2000}, @@ -403,7 +405,7 @@ def test_single_element_sequence_stored_as_scalar(self, simple_model): @pytest.mark.parametrize( "precond_cls,precond_kwargs,precond_name", PRECOND_CONFIGS, - ids=["VP", "VE", "iDDPM", "EDM"], + ids=["TrigFlow", "VP", "VE", "iDDPM", "EDM"], ) class TestNonRegression: """Non-regression tests parameterized across all preconditioner types.""" @@ -611,7 +613,7 @@ def create_fn(): @pytest.mark.parametrize( "precond_cls,precond_kwargs,precond_name", PRECOND_CONFIGS, - ids=["VP", "VE", "iDDPM", "EDM"], + ids=["TrigFlow", "VP", "VE", "iDDPM", "EDM"], ) class TestAllPreconditioners: """Tests that apply to all preconditioner types.""" diff --git a/test/diffusion/test_trigflow.py b/test/diffusion/test_trigflow.py new file mode 100644 index 0000000000..8a32f9acee --- /dev/null +++ b/test/diffusion/test_trigflow.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Check continuous inverse SNR and TrigFlow coordinates.""" + +import pytest +import torch + +from physicsnemo.diffusion import ( + DiffusionExtraTimes, + preconditioners, +) +from physicsnemo.diffusion import noise_schedulers as schedules + +from .helpers import StructuredDiffusionModel + +CONFIGS = ( + ("EDMNoiseScheduler", {}), + ("EDMLogUniformNoiseScheduler", {}), + ("VENoiseScheduler", {}), + ("VPNoiseScheduler", {}), + ("VPNoiseScheduler", {"beta_d": 0.0}), + ("IDDPMNoiseScheduler", {"M": 32}), + ("RectifiedFlowNoiseScheduler", {}), + ("TrigFlowNoiseScheduler", {"sigma_data": 0.5}), + ("TrigFlowNoiseScheduler", {"sigma_data": 1.4}), +) +IDS = ( + "edm", + "edm-uniform", + "ve", + "vp", + "vp-constant", + "iddpm", + "rf", + "trig", + "trig-scaled", +) + + +@pytest.mark.parametrize("name,kwargs", CONFIGS, ids=IDS) +class TestConsistency: + def test_snr_inverse(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.1, 0.4, 0.8], dtype=torch.float64, device=device) + recovered, tangent = torch.func.jvp( + lambda t: scheduler.snr_inv(scheduler.snr(t)), + (t,), + (torch.ones_like(t),), + ) + torch.testing.assert_close(recovered, t) + torch.testing.assert_close(tangent, torch.ones_like(t)) + + def test_coefficient_derivatives(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.1, 0.4, 0.8], dtype=torch.float64, device=device) + for coefficient, derivative in ( + (scheduler.alpha, scheduler.alpha_dot), + (scheduler.sigma, scheduler.sigma_dot), + ): + _, actual = torch.func.jvp(coefficient, (t,), (torch.ones_like(t),)) + torch.testing.assert_close(actual, derivative(t)) + + +@pytest.mark.parametrize("sigma_data", (0.5, 1.4)) +class TestGradientFlow: + def test_trigflow(self, sigma_data): + scheduler = schedules.TrigFlowNoiseScheduler(sigma_data=sigma_data) + model = preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data=sigma_data + ).double() + x = torch.randn(2, 3, dtype=torch.float64, requires_grad=True) + t = torch.tensor([0.4, 0.8], dtype=torch.float64, requires_grad=True) + + def forward(x, t): + prediction, extras = model( + x, (t, DiffusionExtraTimes(r=t / 2, batch_size=[2])) + ) + return ( + prediction, + extras.logvar, + scheduler.snr_inv(scheduler.snr(t)), + ) + + assert torch.autograd.gradcheck(forward, (x, t)) + + +class TestNumerics: + @pytest.mark.parametrize( + "dtype,ratio", + ( + (torch.float16, 0.003), + (torch.float32, 1e-20), + (torch.float64, 1e-200), + ), + ) + def test_vp_small_ratio(self, dtype, ratio, device): + scheduler = schedules.VPNoiseScheduler() + q = torch.tensor([ratio], dtype=dtype, device=device, requires_grad=True) + result = scheduler.snr_inv(q) + (gradient,) = torch.autograd.grad(result.sum(), (q,)) + q64 = q.detach().double() + a = torch.logaddexp(torch.zeros_like(q64), -2 * q64.log()) + root = (scheduler.beta_min**2 + 2 * scheduler.beta_d * a).sqrt() + expected = 2 * a / (scheduler.beta_min + root) + expected_gradient = -2 / (q64 * (1 + q64.square()) * root) + torch.testing.assert_close(result, expected.to(dtype)) + torch.testing.assert_close(gradient, expected_gradient.to(dtype)) + + +@pytest.mark.parametrize("sigma_data", (0.5, 1.4)) +class TestConstructor: + def test_trigflow(self, sigma_data, device): + scheduler = schedules.TrigFlowNoiseScheduler( + sigma_data=sigma_data, t_min=0.1, t_max=1.3 + ) + t = scheduler.sample_time(128, device=device, dtype=torch.float64) + assert torch.all((t >= 0.1) & (t <= 1.3)) + grid = scheduler.timesteps(8, device=device) + assert grid.shape == (9,) + torch.testing.assert_close(grid, torch.linspace(1.3, 0, 9, device=device)) + model = preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data=sigma_data + ).to(device) + x = torch.randn(2, 3, device=device) + t = torch.tensor([0.4, 0.8], device=device) + result, extras = model(x, t) + raw, raw_extras = model.model(x / sigma_data, t) + expected = t.cos()[:, None] * x - sigma_data * t.sin()[:, None] * raw + torch.testing.assert_close(result, expected) + torch.testing.assert_close(extras.logvar, raw_extras.logvar) + torch.testing.assert_close( + scheduler.loss_weight(t), 1 / (sigma_data * t.sin()) ** 2 + ) + + def test_zero_time(self, sigma_data): + scheduler = schedules.TrigFlowNoiseScheduler(sigma_data=sigma_data) + t = torch.zeros(2) + torch.testing.assert_close(scheduler.alpha(t), torch.ones_like(t)) + torch.testing.assert_close(scheduler.sigma(t), t) + torch.testing.assert_close(scheduler.alpha_dot(t), t) + torch.testing.assert_close( + scheduler.sigma_dot(t), torch.full_like(t, sigma_data) + ) + + +@pytest.mark.parametrize("name,kwargs", CONFIGS, ids=IDS) +class TestCompile: + def test_inverse(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.2, 0.7], device=device, requires_grad=True) + fn = torch.compile( + lambda t: scheduler.snr_inv(scheduler.snr(t)), + backend="aot_eager", + fullgraph=True, + ) + result = fn(t) + with torch._dynamo.config.patch(error_on_recompile=True): + torch.testing.assert_close(fn(t), result) + torch.testing.assert_close(result, t) + (gradient,) = torch.autograd.grad(result.sum(), (t,)) + torch.testing.assert_close(gradient, torch.ones_like(t)) + + +class TestDiagnostics: + @pytest.mark.parametrize( + "kwargs", + ( + {"sigma_data": 0}, + {"sigma_data": float("nan")}, + {"t_min": -0.1}, + {"t_max": 1.6}, + {"t_min": 0.7, "t_max": 0.5}, + ), + ) + def test_invalid_scheduler(self, kwargs): + with pytest.raises(ValueError): + schedules.TrigFlowNoiseScheduler(**kwargs) + + @pytest.mark.parametrize("sigma_data", (0, -1, float("nan"))) + def test_invalid_preconditioner(self, sigma_data): + with pytest.raises(ValueError): + preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data + )