TransformerModel.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/Biquad.h"
#include "../Core/DenormalGuard.h"
#include "../Core/DspMath.h"
#include "../Core/Hysteresis.h"
#include "../Core/Oversampling.h"
#include "../Core/StateBlob.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <numbers>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class TransformerModel
{
public:
    // -- Lifecycle ---------------------------------------------------------------
 
    void prepare(const AudioSpec& spec)
    {
        if (!spec.isValid()) return;
        prepared_.store(false, std::memory_order_relaxed);
        spec_ = spec;
        // Everything but the mix runs at the internal rate: the active
        // oversampling factor times the base rate.
        sampleRate_ = static_cast<double>(osFactor_) * spec.sampleRate;
        mixMaxStep_ = static_cast<T>(1.0 / std::max(1.0, spec.sampleRate * 0.02));
        numChannels_ = spec.numChannels;
        maxBlock_ = std::max(spec.maxBlockSize, 1);
 
        if (osFactor_ > 1)
        {
            oversampler_ = std::make_unique<Oversampling<T>>(
                osFactor_, Oversampling<T>::Quality::High);
            oversampler_->prepare(spec);
        }
        else
        {
            oversampler_.reset();
        }
        latency_ = oversampler_ ? oversampler_->getLatency() : 0;
        drySize_ = 1;
        while (drySize_ < latency_ + maxBlock_ + 1) drySize_ <<= 1;
        dryRing_.assign(static_cast<size_t>(numChannels_),
                        std::vector<T>(static_cast<size_t>(drySize_), T(0)));
        dryPos_ = 0;
 
        channels_.assign(static_cast<size_t>(numChannels_), {});
        for (auto& ch : channels_)
        {
            ch.hyst.prepare(sampleRate_);
            // Unbiased iron: wider loop than biased tape (low reversible c).
            // Constant material - set once here, not on every recompute().
            ch.hyst.setParameters(3.5e5, 2.2e4, 1.6e-3, 3.2e4, 0.25);
        }
 
        prepared_.store(true, std::memory_order_relaxed);
        dirty_.store(true, std::memory_order_release);
        reset();
    }
 
    void reset() noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
        for (auto& ch : channels_)
        {
            ch.hyst.reset();
            ch.flux = 0.0;
            ch.vPrev = 0.0;
            ch.vPrev2 = 0.0;
            ch.mPrev = 0.0;
            ch.hpX = ch.hpY = 0.0;
            ch.bell = {};
        }
        for (auto& d : dryRing_)
            std::fill(d.begin(), d.end(), T(0));
        dryPos_ = 0;
        if (oversampler_) oversampler_->reset();
        // Seed the anti-zipper ramps at their targets: no fade-in on start.
        hScaleSm_ = -1.0;
        mScaleSm_ = -1.0;
        currentMix_ = mix_.load(std::memory_order_relaxed);
    }
 
    // -- Parameters (thread-safe) ---------------------------------------------------
 
    void setDrive(T db) noexcept
    {
        if (!std::isfinite(db)) return;
        driveDb_.store(std::clamp(db, T(-12), T(24)), std::memory_order_relaxed);
        dirty_.store(true, std::memory_order_release);
    }
 
    void setCoreSize(T size) noexcept
    {
        if (!std::isfinite(size)) return;
        coreSize_.store(std::clamp(size, T(0), T(1)), std::memory_order_relaxed);
        dirty_.store(true, std::memory_order_release);
    }
 
    void setResonance(T amount) noexcept
    {
        if (!std::isfinite(amount)) return;
        resonance_.store(std::clamp(amount, T(0), T(1)), std::memory_order_relaxed);
        dirty_.store(true, std::memory_order_release);
    }
 
    void setMix(T mix) noexcept
    {
        if (!std::isfinite(mix)) return;
        mix_.store(std::clamp(mix, T(0), T(1)), std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getDrive() const noexcept { return driveDb_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getCoreSize() const noexcept { return coreSize_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getResonance() const noexcept { return resonance_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getMix() const noexcept { return mix_.load(std::memory_order_relaxed); }
 
    void setOversampling(int factor)
    {
        if (factor < 1 || factor > 16 || (factor & (factor - 1)) != 0) return;
        if (factor == osFactor_) return;
        osFactor_ = factor;
        if (prepared_.load(std::memory_order_relaxed))
            prepare(spec_);
    }
 
    [[nodiscard]] int getOversamplingFactor() const noexcept { return osFactor_; }
 
    [[nodiscard]] int getLatency() const noexcept { return latency_; }
 
    [[nodiscard]] int getLatencySamples() const noexcept { return getLatency(); }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        StateWriter w(stateId("XFMR"), 2);
        // Explicit float casts: the blob stores float, and with T = double the
        // unqualified write(key, double) would be ambiguous (float/int32/bool).
        w.write("drive", static_cast<float>(driveDb_.load(std::memory_order_relaxed)));
        w.write("coreSize", static_cast<float>(coreSize_.load(std::memory_order_relaxed)));
        w.write("resonance", static_cast<float>(resonance_.load(std::memory_order_relaxed)));
        w.write("mix", static_cast<float>(mix_.load(std::memory_order_relaxed)));
        w.write("oversampling", osFactor_);
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("XFMR")) return false;
        setDrive(static_cast<T>(r.read("drive", 0.0f)));
        setCoreSize(static_cast<T>(r.read("coreSize", 0.5f)));
        setResonance(static_cast<T>(r.read("resonance", 0.3f)));
        setMix(static_cast<T>(r.read("mix", 1.0f)));
        setOversampling(r.read("oversampling", 2));
        return true;
    }
 
    // -- Processing -------------------------------------------------------------------
 
    void processBlock(AudioBufferView<T> buffer) noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
        DenormalGuard guard;
 
        const int nCh = std::min(buffer.getNumChannels(), numChannels_);
        const int nS = buffer.getNumSamples();
        if (nCh == 0 || nS == 0) return;
 
        // Acquire pairs with the setters' release stores so the recompute
        // always sees the values published before the flag.
        if (dirty_.load(std::memory_order_relaxed)
            && dirty_.exchange(false, std::memory_order_acquire))
            recompute();
 
        for (int off = 0; off < nS; off += maxBlock_)
            processChunk(buffer.getSubView(off, std::min(maxBlock_, nS - off)), nCh);
    }
 
private:
    void processChunk(AudioBufferView<T> buffer, int nCh) noexcept
    {
        const int nS = buffer.getNumSamples();
 
        // Non-finite guard: a NaN/Inf input would poison the leaky
        // integrator (flux/vPrev), the algebraic inverse (vPrev2/mPrev), the
        // DC-blocker (hpX/hpY), the bell and the oversampler permanently.
        // Replace the bad sample with silence (dry too) so a transient glitch
        // cannot kill the channel for the rest of the stream.
        for (int ch = 0; ch < nCh; ++ch)
        {
            T* d = buffer.getChannel(ch);
            for (int i = 0; i < nS; ++i)
                if (!std::isfinite(d[i])) d[i] = T(0);
        }
 
        // Rate-limited mix ramp (moveTowards, exact landing; settled it
        // reduces to the constant, bit-identically). A per-block ramp landed
        // in 0.7 ms with 32-sample blocks.
        // A hard flip on the differentiated wet stream clicked at 4.6x the
        // steady-state sample delta.
        const T mixTarget = mix_.load(std::memory_order_relaxed);
        const T mixStart  = currentMix_;
 
        // Dry snapshot, read back delayed to the oversampler's latency.
        for (int ch = 0; ch < nCh; ++ch)
        {
            const T* in = buffer.getChannel(ch);
            auto& dry = dryRing_[static_cast<size_t>(ch)];
            int dp = dryPos_;
            for (int i = 0; i < nS; ++i)
            {
                dry[static_cast<size_t>(dp)] = in[i];
                dp = (dp + 1) & (drySize_ - 1);
            }
        }
 
        const bool osOn = (oversampler_ != nullptr);
        auto osView = osOn ? oversampler_->upsample(buffer) : buffer;
        const int osN = osView.getNumSamples();
 
        // Anti-zipper: GEOMETRIC in-block ramps toward the recompute()
        // targets (~50 ms across blocks), shared by all channels. Two
        // reasons: this model differentiates its output (x 2*fs), so a flat
        // per-block gain step clicks; and (hScale, mScale) is a calibrated
        // compensation pair spanning orders of magnitude - interpolating it
        // LINEARLY passes through wildly over-gained states (+24 dB
        // measured mid-drag), while power-law interpolation keeps the pair
        // on the calibration curve m ~ C/h^g throughout the transition.
        if (hScaleSm_ <= 0.0) { hScaleSm_ = hScale_; mScaleSm_ = mScale_; }
        const double kSm = 1.0 - std::exp(-static_cast<double>(osN) / (0.050 * sampleRate_));
        const double hEnd = hScaleSm_ * std::pow(hScale_ / hScaleSm_, kSm);
        const double mEnd = mScaleSm_ * std::pow(mScale_ / mScaleSm_, kSm);
        const double hRat = std::pow(hEnd / hScaleSm_, 1.0 / static_cast<double>(osN));
        const double mRat = std::pow(mEnd / mScaleSm_, 1.0 / static_cast<double>(osN));
 
        for (int ch = 0; ch < nCh; ++ch)
        {
            T* d = osView.getChannel(ch);
            auto& st = channels_[static_cast<size_t>(ch)];
            double hSm = hScaleSm_, mSm = mScaleSm_;
 
            for (int i = 0; i < osN; ++i)
            {
                hSm *= hRat;
                mSm *= mRat;
                const double x = static_cast<double>(d[i]);
 
                // Leaky trapezoidal integrator: winding voltage -> flux.
                const double fluxNew = leak_ * st.flux + halfT_ * (x + st.vPrev);
                st.vPrev = x;
 
                // Core hysteresis on the flux (JA, unbiased-iron parameters).
                const double m = mSm * static_cast<double>(
                    st.hyst.processSample(static_cast<T>(hSm * fluxNew)));
 
                // Algebraic inverse of the integrator (flux -> voltage). The
                // damped feedback pole (rho < 1) tames z = -1: hysteresis
                // distortion is new content the inverse never saw, and an
                // undamped differentiator would ring at Nyquist forever. The
                // leak alpha still cancels exactly; the only linear residue
                // is (1+z^-1)/(1+rho z^-1), transparent to within 0.15 dB.
                double v = (m - leak_ * st.mPrev) * invHalfT_ - kDiffRho * st.vPrev2;
                st.vPrev2 = v;
                st.mPrev = m;
                st.flux = fluxNew;
 
                // Magnetizing-inductance corner (one-pole high-pass).
                const double hp = hpA_ * (st.hpY + v - st.hpX);
                st.hpX = v;
                st.hpY = hp;
 
                // Leakage/capacitance bell.
                d[i] = static_cast<T>(st.bell.process(hp));
            }
        }
        hScaleSm_ = hEnd;
        mScaleSm_ = mEnd;
        if (osOn) oversampler_->downsample(buffer);
 
        for (int ch = 0; ch < nCh; ++ch)
        {
            T* d = buffer.getChannel(ch);
            const auto& dry = dryRing_[static_cast<size_t>(ch)];
            for (int i = 0; i < nS; ++i)
            {
                const int idx = (dryPos_ + i - latency_) & (drySize_ - 1);
                const T drySample = dry[static_cast<size_t>(idx)];
                const T mixVal = moveTowards(mixStart, mixTarget, mixMaxStep_ * static_cast<T>(i + 1));
                d[i] = drySample + (d[i] - drySample) * mixVal;
            }
        }
        currentMix_ = moveTowards(mixStart, mixTarget, mixMaxStep_ * static_cast<T>(nS));
        dryPos_ = (dryPos_ + nS) & (drySize_ - 1);
    }
 
    static constexpr double kDiffRho = 0.974;   
 
    struct BellSection
    {
        double b0 = 1.0, b1 = 0.0, b2 = 0.0, a1 = 0.0, a2 = 0.0;
        double z1 = 0.0, z2 = 0.0;
 
        [[nodiscard]] double process(double x) noexcept
        {
            const double y = b0 * x + z1;
            z1 = b1 * x - a1 * y + z2;
            z2 = b2 * x - a2 * y;
            return y;
        }
    };
 
    struct ChannelState
    {
        Hysteresis<T> hyst;
        double flux = 0.0;
        double vPrev = 0.0;     
        double vPrev2 = 0.0;    
        double mPrev = 0.0;
        double hpX = 0.0, hpY = 0.0;
        BellSection bell;
    };
 
    void recompute() noexcept
    {
        const double drive = std::pow(10.0, static_cast<double>(
            driveDb_.load(std::memory_order_relaxed)) / 20.0);
        const double size = static_cast<double>(coreSize_.load(std::memory_order_relaxed));
        const double res = static_cast<double>(resonance_.load(std::memory_order_relaxed));
 
        // The integrator leak keeps the flux bounded; its exact algebraic
        // inverse cancels it, so the audible LF rolloff comes only from the
        // explicit magnetizing-corner high-pass below.
        const double cornerHz = 40.0 * std::pow(0.125, size);   // 40 -> 5 Hz
        leak_ = std::exp(-2.0 * std::numbers::pi * cornerHz / sampleRate_);
        hpA_ = 1.0 - 2.0 * std::numbers::pi * cornerHz / sampleRate_;
        halfT_ = 0.5 / sampleRate_;
        invHalfT_ = 2.0 * sampleRate_;
 
        // Flux scale: 0 dBFS at 30 Hz reaches H = 1.1a at nominal drive and
        // nominal core; bigger cores take proportionally more flux.
        const double fluxRef = 1.0 / (2.0 * std::numbers::pi * 30.0);
        const double headroom = 0.6 + 0.8 * size;
        hScale_ = drive * 1.1 * 2.2e4 / (fluxRef * headroom);
 
        // Loudness calibration at PROGRAM level: 100 Hz at -12 dBFS, the
        // flux regime music actually drives the core into (flux ~ V/f, so
        // lows dominate). Calibrating small-signal (the old 0.05 @ 1 kHz)
        // probed the JA virgin curve, whose susceptibility is an order of
        // magnitude below the major loop's mean slope - at high drive the
        // output then came out ~+21 dB over the input (measured x11),
        // heard as violent level jumps while dragging the drive slider.
        {
            Hysteresis<T> cal;
            cal.prepare(sampleRate_);
            cal.setParameters(3.5e5, 2.2e4, 1.6e-3, 3.2e4, 0.25);
            double flux = 0.0, vPrev = 0.0, mPrev = 0.0, vPrev2 = 0.0;
            double inSq = 0.0, outSq = 0.0;
            const int n = static_cast<int>(0.04 * sampleRate_);   // 2+2 cycles of 100 Hz
            for (int i = 0; i < n; ++i)
            {
                const double x = 0.25 * std::sin(2.0 * std::numbers::pi * 100.0 * i / sampleRate_);
                const double fluxNew = leak_ * flux + halfT_ * (x + vPrev);
                vPrev = x;
                const double m = static_cast<double>(
                    cal.processSample(static_cast<T>(hScale_ * fluxNew)));
                double v = (m - leak_ * mPrev) * invHalfT_ - kDiffRho * vPrev2;
                vPrev2 = v;
                mPrev = m;
                flux = fluxNew;
                if (i >= n / 2)
                {
                    inSq += x * x;
                    outSq += v * v;
                }
            }
            mScale_ = (outSq > 0.0) ? std::sqrt(inSq / outSq) : 1.0;
        }
 
        // HF bell: Jensen-style leakage resonance mapped into the top octave.
        const double bellHz = std::min(12000.0 + 6000.0 * res, 0.42 * sampleRate_);
        const double bellDb = 2.5 * res;
        const auto bc = BiquadCoeffs::makePeak(sampleRate_, bellHz, 0.8, bellDb);
        for (auto& ch : channels_)
        {
            const double pz1 = ch.bell.z1, pz2 = ch.bell.z2;
            ch.bell = BellSection { bc.b0, bc.b1, bc.b2, bc.a1, bc.a2, 0.0, 0.0 };
            ch.bell.z1 = pz1;
            ch.bell.z2 = pz2;
        }
    }
 
    // -- Members --------------------------------------------------------------------
    AudioSpec spec_ {};         
    double sampleRate_ = 48000.0;   
    int numChannels_ = 0;
    int maxBlock_ = 1;
    int osFactor_ = 2;          
    int latency_ = 0;
    std::unique_ptr<Oversampling<T>> oversampler_;
    std::vector<std::vector<T>> dryRing_;   
    int drySize_ = 1;
    int dryPos_ = 0;
    std::atomic<bool> prepared_ { false };
 
    std::vector<ChannelState> channels_;
 
    double leak_ = 0.999;
    double hpA_ = 0.999;
    double halfT_ = 0.5 / 48000.0;
    double invHalfT_ = 96000.0;
    double hScale_ = 1.0;
    double mScale_ = 1.0;
    double hScaleSm_ = -1.0;    
    double mScaleSm_ = -1.0;    
    T currentMix_ = T(1);       
    T mixMaxStep_ = T(1.0 / 960.0); 
 
    std::atomic<T> driveDb_ { T(0) };
    std::atomic<T> coreSize_ { T(0.5) };
    std::atomic<T> resonance_ { T(0.3) };
    std::atomic<T> mix_ { T(1) };
    std::atomic<bool> dirty_ { true };
};
 
} // namespace dspark