Search the stereo angle at high complexity

internal/celt/encode_bands.go

@@ -11,6 +11,10 @@ import (
 	"github.com/pion/opus/internal/slicetools"
 )
 
+// thetaRDOComplexity is the encoder complexity at which libopus turns on the
+// rate-distortion search over the stereo angle (celt/bands.c).
+const thetaRDOComplexity = 8
+
 type bandEncodeState struct {
 	rangeEncoder   *rangecoding.Encoder
 	seed           uint32
@@ -22,6 +26,45 @@ type bandEncodeState struct {
 	// intensityStereo weighs its downmix by. libopus hands quant_all_bands the
 	// unquantized bandE from compute_band_energies, so this mirrors that.
 	bandEnergy [2][maxBands]float32
+	// thetaRDO turns on the rate-distortion search over the stereo angle
+	// (libopus gates it on complexity>=8). rdoScratch holds the copies the
+	// search needs to rewind between the two candidates.
+	thetaRDO   bool
+	rdoScratch *[4][]float32
+	rdoState   *[2]rangecoding.State
+}
+
+// rdoSave copies a band's two channels aside so a speculative encode can be
+// undone. slot 0 holds the state before the first candidate, slot 1 the result
+// of the first candidate.
+func (s *bandEncodeState) rdoSave(slot int, x, y []float32) {
+	s.rdoScratch[2*slot] = append(s.rdoScratch[2*slot][:0], x...)
+	s.rdoScratch[2*slot+1] = append(s.rdoScratch[2*slot+1][:0], y...)
+}
+
+func (s *bandEncodeState) rdoRestore(slot int, x, y []float32) {
+	copy(x, s.rdoScratch[2*slot])
+	copy(y, s.rdoScratch[2*slot+1])
+}
+
+// channelWeights mirrors libopus compute_channel_weights: the quieter channel
+// is nudged up so a near-silent side does not dominate the distortion.
+func channelWeights(ex, ey float32) (float32, float32) {
+	minE := min32(ex, ey)
+
+	return ex + minE/3, ey + minE/3
+}
+
+// bandDistortion is the weighted correlation between the original band and its
+// reconstruction. Higher is better, which is the direction libopus compares.
+func bandDistortion(xOrig, yOrig, x, y []float32, wx, wy float32) float64 {
+	var dx, dy float64
+	for i := range x {
+		dx += float64(xOrig[i]) * float64(x[i])
+		dy += float64(yOrig[i]) * float64(y[i])
+	}
+
+	return float64(wx)*dx + float64(wy)*dy
 }
 
 func (s *bandEncodeState) floatScratch(n int) []float32 {
@@ -542,6 +585,7 @@ func quantBandStereo(
 	yScratch [2][]int,
 	absXScratch, signScratch [2][]float32,
 	cwrsScratch []uint32,
+	thetaRound int,
 ) uint {
 	if n == 1 {
 		xSign := uint32(0)
@@ -587,7 +631,7 @@ func quantBandStereo(
 	itheta := 0
 	invert := false
 	if qn != 1 {
-		thetaSym = quantizeStereoBandTheta(x, y, qn)
+		thetaSym = quantizeStereoBandTheta(x, y, qn, thetaRound)
 		encodeBandTheta(thetaSym, qn, n, true, blocks, state.rangeEncoder)
 		itheta = thetaSym * 16384 / qn
 		// libopus compute_theta (celt/bands.c:866-871): a zero angle means the
@@ -754,7 +798,11 @@ func stereoSplit(x []float32, y []float32, n int) {
 	}
 }
 
-func quantizeStereoBandTheta(x []float32, y []float32, qn int) int {
+// quantizeStereoBandTheta returns the coded angle symbol. thetaRound is 0 for
+// the plain nearest-value choice; -1 and +1 ask for the two candidates the
+// rate-distortion search compares, biased toward the ends of the range so the
+// pair straddles the unrounded angle (libopus celt/bands.c compute_theta).
+func quantizeStereoBandTheta(x []float32, y []float32, qn, thetaRound int) int {
 	if qn <= 1 {
 		return 0
 	}
@@ -774,9 +822,23 @@ func quantizeStereoBandTheta(x []float32, y []float32, qn int) int {
 	}
 
 	theta := math.Atan2(math.Sqrt(ey), math.Sqrt(ex))
-	symbol := int(math.Round(theta * float64(qn) / (0.5 * math.Pi)))
+	// The reference works in a Q14 angle; keeping the same domain here makes
+	// the bias below the same integer it uses.
+	raw := theta * 16384 / (0.5 * math.Pi)
+	if thetaRound == 0 {
+		return min(qn, max(0, int(math.Round(raw*float64(qn)/16384))))
+	}
 
-	return min(qn, max(0, symbol))
+	bias := 32767 / qn
+	if raw <= 8192 {
+		bias = -bias
+	}
+	down := min(qn-1, max(0, int(math.Floor((raw*float64(qn)+float64(bias))/16384))))
+	if thetaRound < 0 {
+		return down
+	}
+
+	return down + 1
 }
 
 func quantAllBandsStereo(
@@ -918,25 +980,22 @@ func quantAllBandsStereo(
 				yScratch[1], absXScratch[1], signScratch[1], cwrsScratch,
 			)
 		} else {
-			xMask = quantBandStereo(
-				band,
-				x[bandStart:bandEnd],
-				y[bandStart:bandEnd],
-				bandWidth,
-				bandBits,
-				info.spread,
-				blocks,
-				info.allocation.intensity,
-				info.tfChange[band],
-				lowband,
-				&remainingBits,
-				info.lm,
-				1,
-				lowbandScratch,
-				xMask|yMask,
-				state,
-				yScratch, absXScratch, signScratch, cwrsScratch,
-			)
+			bx := x[bandStart:bandEnd]
+			by := y[bandStart:bandEnd]
+			quant := func(round int, mask uint) uint {
+				return quantBandStereo(
+					band, bx, by, bandWidth, bandBits, info.spread, blocks,
+					info.allocation.intensity, info.tfChange[band], lowband,
+					&remainingBits, info.lm, 1, lowbandScratch, mask, state,
+					yScratch, absXScratch, signScratch, cwrsScratch, round,
+				)
+			}
+			if state.thetaRDO && band < info.allocation.intensity {
+				xMask = quantBandStereoRDO(state, quant, bx, by, xMask|yMask, &remainingBits,
+					state.bandEnergy[0][band], state.bandEnergy[1][band])
+			} else {
+				xMask = quant(0, xMask|yMask)
+			}
 			yMask = xMask
 		}
 
@@ -950,3 +1009,41 @@ func quantAllBandsStereo(
 
 	return collapseMasks
 }
+
+// quantBandStereoRDO encodes the band twice — rounding the stereo angle down
+// and then up — and keeps whichever reconstruction tracks the original more
+// closely. The two encodes are speculative, so the range coder, the band and
+// the noise-fill seed all rewind between them (libopus celt/bands.c).
+func quantBandStereoRDO(
+	state *bandEncodeState, quant func(round int, mask uint) uint,
+	x, y []float32, mask uint, remainingBits *int, ex, ey float32,
+) uint {
+	wx, wy := channelWeights(ex, ey)
+
+	state.rangeEncoder.SaveInto(&state.rdoState[0])
+	seedBefore, bitsBefore := state.seed, *remainingBits
+	state.rdoSave(0, x, y)
+
+	maskDown := quant(-1, mask)
+	distDown := bandDistortion(state.rdoScratch[0], state.rdoScratch[1], x, y, wx, wy)
+
+	state.rangeEncoder.SaveInto(&state.rdoState[1])
+	seedDown, bitsDown := state.seed, *remainingBits
+	state.rdoSave(1, x, y)
+
+	state.rangeEncoder.Restore(&state.rdoState[0])
+	state.seed, *remainingBits = seedBefore, bitsBefore
+	state.rdoRestore(0, x, y)
+
+	maskUp := quant(1, mask)
+	distUp := bandDistortion(state.rdoScratch[0], state.rdoScratch[1], x, y, wx, wy)
+	if distDown < distUp {
+		return maskUp
+	}
+
+	state.rangeEncoder.Restore(&state.rdoState[1])
+	state.seed, *remainingBits = seedDown, bitsDown
+	state.rdoRestore(1, x, y)
+
+	return maskDown
+}

internal/celt/encoder.go

@@ -47,6 +47,10 @@ type Encoder struct {
 	normalizedBands   [2][]float32
 	pitchBuf          []float32
 	pitchChannels     [2][]float32
+	// rdoScratch and rdoState back the stereo-angle search. They live here so
+	// the per-frame band state can borrow them instead of allocating.
+	rdoScratch [4][]float32
+	rdoState   [2]rangecoding.State
 
 	spreadAverage      int
 	hfAverage          int
@@ -790,6 +794,33 @@ func (e *Encoder) updateVBRReservoir(vbrRate, rawTarget, roundedTarget int) {
 	}
 }
 
+// newBandState wires up the per-frame state the band quantiser works from,
+// including the linear band energies libopus hands quant_all_bands.
+func (e *Encoder) newBandState(
+	info *frameSideInfo, logBandAmp [2][maxBands]float32,
+) bandEncodeState {
+	state := bandEncodeState{
+		rangeEncoder:   &e.rangeEncoder,
+		seed:           e.rng,
+		norm:           e.bandNorm[:0],
+		lowbandScratch: e.bandLowScratch[:0],
+		collapseMasks:  e.bandCollapseMasks[:0],
+		// libopus only searches the stereo angle at complexity 8 and up: it
+		// encodes every stereo band twice, so the cost is real.
+		thetaRDO:   e.complexity >= thetaRDOComplexity && info.channelCount == 2,
+		rdoScratch: &e.rdoScratch,
+		rdoState:   &e.rdoState,
+	}
+	for ch := range info.channelCount {
+		for band := info.startBand; band < info.endBand; band++ {
+			state.bandEnergy[ch][band] = float32(math.Pow(2,
+				float64(logBandAmp[ch][band]+energyMeans[band])))
+		}
+	}
+
+	return state
+}
+
 // EncodeFrame encodes one CELT frame from float PCM into dst.
 // It returns the number of bytes written. dst must be at least frameBytes long.
 //
@@ -924,19 +955,7 @@ func (e *Encoder) EncodeFrame(pcm [][]float32, dst []byte, frameBytes, startBand
 	e.encodeFineEnergy(&info, info.allocation.fineQuant, targetLogE)
 
 	totalBits := (int(info.totalBits) << bitResolution) - info.antiCollapseRsv
-	bandState := bandEncodeState{
-		rangeEncoder:   &e.rangeEncoder,
-		seed:           e.rng,
-		norm:           e.bandNorm[:0],
-		lowbandScratch: e.bandLowScratch[:0],
-		collapseMasks:  e.bandCollapseMasks[:0],
-	}
-	for ch := range info.channelCount {
-		for band := info.startBand; band < info.endBand; band++ {
-			bandState.bandEnergy[ch][band] = float32(math.Pow(2,
-				float64(analysis.logBandAmp[ch][band]+energyMeans[band])))
-		}
-	}
+	bandState := e.newBandState(&info, analysis.logBandAmp)
 	shape0 := normalized[0]
 	if info.channelCount == 2 {
 		shape1 := normalized[1]

internal/celt/encoder_test.go

@@ -220,6 +220,7 @@ func TestQuantBandStereoN1(t *testing.T) {
 		[2][]float32{make([]float32, 1), make([]float32, 1)},
 		[2][]float32{make([]float32, 1), make([]float32, 1)},
 		make([]uint32, cwrsMaxPulseCount+2),
+		0,
 	)
 	assert.Equal(t, uint(1), mask)
 }
@@ -241,6 +242,7 @@ func TestQuantBandStereoN2(t *testing.T) {
 		[2][]float32{make([]float32, 2), make([]float32, 2)},
 		[2][]float32{make([]float32, 2), make([]float32, 2)},
 		make([]uint32, cwrsMaxPulseCount+2),
+		0,
 	)
 	assert.Greater(t, enc.rangeEncoder.FinalRange(), uint32(0))
 }

internal/celt/theta_rdo_test.go

@@ -0,0 +1,111 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package celt
+
+import (
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestThetaRDORoundTrip drives the rate-distortion search, which encodes every
+// stereo band twice and rewinds the range coder, the band and the noise-fill
+// seed in between. An incomplete rewind desynchronises the range coder, so the
+// final-range check is what actually guards the search.
+func TestThetaRDORoundTrip(t *testing.T) {
+	encoder := NewEncoder()
+	encoder.SetComplexity(thetaRDOComplexity)
+	decoder := NewDecoder()
+
+	frameSampleCount := shortBlockSampleCount << maxLM
+	left := make([]float32, frameSampleCount)
+	right := make([]float32, frameSampleCount)
+	phase := 0.0
+	out := make([]float32, frameSampleCount*2)
+	for frame := range 5 {
+		for i := range frameSampleCount {
+			// Both channels move so the angle lands somewhere different every
+			// frame; a fixed pair would only ever exercise one rounding.
+			left[i] = float32(math.Sin(phase) + 0.5*math.Sin(phase*3.7))
+			right[i] = float32(math.Sin(phase*1.03+0.4) - 0.3*math.Sin(phase*7.1))
+			phase += 2 * math.Pi * 500 / sampleRate
+		}
+		data := encodeFrame(t, &encoder, [][]float32{left, right}, 240)
+		require.NotEmptyf(t, data, "frame %d", frame)
+		require.NoErrorf(t, decoder.Decode(data, out, true, 2, frameSampleCount, 0, maxBands),
+			"frame %d", frame)
+		assert.Equalf(t, encoder.FinalRange(), decoder.FinalRange(),
+			"range coder out of sync at frame %d", frame)
+	}
+}
+
+// TestThetaRDOChangesTheBitstream pins the gate from the other side: below the
+// threshold the encoder must take the plain path, so the two bitstreams have to
+// differ. Without this the search could silently no-op and every other test
+// here would still pass.
+func TestThetaRDOChangesTheBitstream(t *testing.T) {
+	frameSampleCount := shortBlockSampleCount << maxLM
+	left := make([]float32, frameSampleCount)
+	right := make([]float32, frameSampleCount)
+	for i := range frameSampleCount {
+		left[i] = float32(math.Sin(2 * math.Pi * 440 * float64(i) / sampleRate))
+		right[i] = float32(math.Sin(2 * math.Pi * 660 * float64(i) / sampleRate))
+	}
+
+	plain := NewEncoder()
+	plain.SetComplexity(thetaRDOComplexity - 1)
+	searched := NewEncoder()
+	searched.SetComplexity(thetaRDOComplexity)
+
+	plainData := encodeFrame(t, &plain, [][]float32{left, right}, 240)
+	searchedData := encodeFrame(t, &searched, [][]float32{left, right}, 240)
+
+	assert.NotEqual(t, plainData, searchedData, "the search should change the chosen angles")
+	assert.NotEqual(t, plain.FinalRange(), searched.FinalRange(),
+		"a different set of angles must land on a different range-coder state")
+}
+
+func TestQuantizeStereoBandThetaRoundingBrackets(t *testing.T) {
+	// The two candidates have to bracket the nearest-value choice: if they
+	// collapsed onto the same symbol the search would compare an encode with
+	// itself and never pick anything.
+	x := []float32{0.9, 0.4, -0.2, 0.7, 0.1, -0.5, 0.3, 0.8}
+	y := []float32{0.2, -0.6, 0.5, 0.1, -0.9, 0.4, -0.3, 0.2}
+
+	for _, qn := range []int{2, 4, 8, 16, 32} {
+		down := quantizeStereoBandTheta(x, y, qn, -1)
+		up := quantizeStereoBandTheta(x, y, qn, 1)
+		nearest := quantizeStereoBandTheta(x, y, qn, 0)
+
+		assert.Equalf(t, down+1, up, "qn=%d: candidates must be adjacent", qn)
+		assert.GreaterOrEqualf(t, nearest, down, "qn=%d", qn)
+		assert.LessOrEqualf(t, nearest, up, "qn=%d", qn)
+		assert.GreaterOrEqualf(t, down, 0, "qn=%d", qn)
+		assert.LessOrEqualf(t, up, qn, "qn=%d", qn)
+	}
+}
+
+func TestChannelWeightsLiftTheQuieterChannel(t *testing.T) {
+	// A near-silent side must not make its own distortion irrelevant, so the
+	// reference nudges both weights up by a third of the smaller energy.
+	wx, wy := channelWeights(9, 3)
+	assert.InDelta(t, 10.0, wx, 1e-6)
+	assert.InDelta(t, 4.0, wy, 1e-6)
+
+	// Equal energies stay equal.
+	wx, wy = channelWeights(6, 6)
+	assert.InDelta(t, wx, wy, 1e-6)
+}
+
+func TestBandDistortionRewardsTheCloserReconstruction(t *testing.T) {
+	orig := []float32{1, 0, -1, 0}
+	exact := []float32{1, 0, -1, 0}
+	poor := []float32{0, 1, 0, -1}
+
+	good := bandDistortion(orig, orig, exact, exact, 1, 1)
+	bad := bandDistortion(orig, orig, poor, poor, 1, 1)
+	assert.Greater(t, good, bad, "a matching reconstruction must score higher")
+}