Add inverse SNR and TrigFlow components

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.

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
 ----------------
 

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`
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

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
 

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"],

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"],

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(

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"],