Feat/stereo saving

internal/celt/allocation.go

@@ -647,7 +647,7 @@ func chooseAllocationTrim(
 	mdct [2][]float32,
 	channelCount, lm, endBand int,
 	totalBits uint,
-	tfEstimate float32,
+	tfEstimate float32, intensity int, stereoSaving *float32,
 ) int {
 	frameSampleCount := shortBlockSampleCount << lm
 	equivRate := int(totalBits) * sampleRate / frameSampleCount
@@ -671,23 +671,25 @@ func chooseAllocationTrim(
 		scale := 1 << lm
 		var corrSum float32
 		for band := 0; band < 8 && band < endBand; band++ {
-			start := scale * int(bandEdges[band])
-			end := scale * int(bandEdges[band+1])
-			var dot, l2, r2 float32
-			for i := start; i < end; i++ {
-				dot += mdct[0][i] * mdct[1][i]
-				l2 += mdct[0][i] * mdct[0][i]
-				r2 += mdct[1][i] * mdct[1][i]
-			}
-			if l2 > 1e-30 && r2 > 1e-30 {
-				corrSum += dot / sqrtf(l2*r2)
-			}
+			corrSum += bandCorrelation(mdct, scale, band)
 		}
 		avgCorr := abs32(corrSum / 8.0)
 		if avgCorr > 1.0 {
 			avgCorr = 1.0
 		}
+		// minXC is the weakest correlation across the intensity-coded range;
+		// a single decorrelated band there means mid/side is not saving much,
+		// however redundant the low bands look.
+		minXC := avgCorr
+		for band := 8; band < intensity && band < endBand; band++ {
+			minXC = min32(minXC, abs32(bandCorrelation(mdct, scale, band)))
+		}
+		if minXC > 1.0 {
+			minXC = 1.0
+		}
 		logXC := float32(math.Log2(1.001 - float64(avgCorr*avgCorr)))
+		logXC2 := max32(0.5*logXC, float32(math.Log2(1.001-float64(minXC*minXC))))
+		*stereoSaving = min32(*stereoSaving+0.25, -0.5*logXC2)
 		trim += max32(-4.0, 0.75*logXC)
 	}
 
@@ -723,6 +725,25 @@ func chooseAllocationTrim(
 	return trimIndex
 }
 
+// bandCorrelation returns the cosine similarity of the two channels over one
+// band, which is what the reference gets from the inner product of the
+// normalised spectrum.
+func bandCorrelation(mdct [2][]float32, scale, band int) float32 {
+	start := scale * int(bandEdges[band])
+	end := scale * int(bandEdges[band+1])
+	var dot, l2, r2 float32
+	for i := start; i < end; i++ {
+		dot += mdct[0][i] * mdct[1][i]
+		l2 += mdct[0][i] * mdct[0][i]
+		r2 += mdct[1][i] * mdct[1][i]
+	}
+	if l2 <= 1e-30 || r2 <= 1e-30 {
+		return 0
+	}
+
+	return dot / sqrtf(l2*r2)
+}
+
 func abs32(x float32) float32 {
 	if x < 0 {
 		return -x

internal/celt/analysis_test.go

@@ -395,7 +395,7 @@ func TestChooseAllocationTrimDefault(t *testing.T) {
 	mdct := makeFlatMDCT()
 	trim := chooseAllocationTrim(
 		[2][maxBands]float32{logBandAmp, logBandAmp},
-		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0,
+		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0, 0, new(float32),
 	)
 	assert.InDelta(t, 5, trim, 1, "flat spectrum at 128kbps should stay near default")
 }
@@ -406,7 +406,7 @@ func TestChooseAllocationTrimLowBitrate(t *testing.T) {
 	mdct := makeFlatMDCT()
 	trim := chooseAllocationTrim(
 		[2][maxBands]float32{logBandAmp, logBandAmp},
-		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 32*8*50, 0,
+		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 32*8*50, 0, 0, new(float32),
 	)
 	assert.LessOrEqual(t, trim, 5, "low bitrate should bias trim downward")
 }
@@ -418,9 +418,11 @@ func TestChooseAllocationTrimSpectralTilt(t *testing.T) {
 	highHeavy := makeTiltedLogBandAmp(+1.0) // bandas altas con más energía
 	mdct := makeFlatMDCT()
 	trimLow := chooseAllocationTrim([2][maxBands]float32{lowHeavy, lowHeavy},
-		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0)
+		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0, 0, new(float32),
+	)
 	trimHigh := chooseAllocationTrim([2][maxBands]float32{highHeavy, highHeavy},
-		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0)
+		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0, 0, new(float32),
+	)
 	assert.Greater(t, trimLow, trimHigh, "low-heavy spectrum should bias trim upward (more bits to lows)")
 }
 
@@ -429,12 +431,14 @@ func TestChooseAllocationTrimStereoCorrelated(t *testing.T) {
 	logBandAmp := makeFlatLogBandAmp(0.0)
 	mdct := makeSineMDCT(440) // mismo contenido en ambos canales
 	trimCorr := chooseAllocationTrim([2][maxBands]float32{logBandAmp, logBandAmp},
-		[2][]float32{mdct, mdct}, 2, maxLM, maxBands, 128*8*50, 0)
+		[2][]float32{mdct, mdct}, 2, maxLM, maxBands, 128*8*50, 0, 0, new(float32),
+	)
 
 	// L y R decorrelated → trim sin ajuste stereo.
 	mdctR := makeNoiseMDCT(42)
 	trimDecorr := chooseAllocationTrim([2][maxBands]float32{logBandAmp, logBandAmp},
-		[2][]float32{mdct, mdctR}, 2, maxLM, maxBands, 128*8*50, 0)
+		[2][]float32{mdct, mdctR}, 2, maxLM, maxBands, 128*8*50, 0, 0, new(float32),
+	)
 
 	assert.Less(t, trimCorr, trimDecorr, "correlated stereo should have lower trim than decorrelated")
 }
@@ -445,9 +449,10 @@ func TestChooseAllocationTrimTFEstimate(t *testing.T) {
 	logBandAmp := makeFlatLogBandAmp(0.0)
 	mdct := makeFlatMDCT()
 	trimFlat := chooseAllocationTrim([2][maxBands]float32{logBandAmp, logBandAmp},
-		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0)
+		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 0, 0, new(float32),
+	)
 	trimTF := chooseAllocationTrim([2][maxBands]float32{logBandAmp, logBandAmp},
-		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 1.0)
+		[2][]float32{mdct, mdct}, 1, maxLM, maxBands, 128*8*50, 1.0, 0, new(float32))
 	assert.Equal(t, trimFlat-2, trimTF, "tf_estimate of 1.0 should drop the trim by 2")
 }
 

internal/celt/encoder.go

@@ -53,6 +53,9 @@ type Encoder struct {
 	tapsetDecision     int
 	prevSpreadDecision int
 	prevIntensityBand  int
+	// stereoSaving estimates how many bits mid/side is saving over plain
+	// stereo; the VBR target spends less when the channels are redundant.
+	stereoSaving float32
 	// lastCodedBands feeds the band-skip hysteresis in computeAllocation
 	// (st->lastCodedBands in libopus celt_encoder.c). Zero means "no previous
 	// frame", which the update below seeds directly instead of clamping.
@@ -134,6 +137,7 @@ func (e *Encoder) Reset() {
 	e.tapsetDecision = 0
 	e.prevSpreadDecision = defaultSpreadDecision
 	e.prevIntensityBand = 0
+	e.stereoSaving = 0
 	e.lastCodedBands = 0
 	e.consecTransient = 0
 	e.analysis.prefilter = postFilterState{}
@@ -359,7 +363,7 @@ func (e *Encoder) encodeDynamicAllocation(info *frameSideInfo, offsets [maxBands
 // enough budget left to signal it.
 func (e *Encoder) encodeAllocationTrim(
 	info *frameSideInfo, logBandAmp [2][maxBands]float32, mdct [2][]float32, totalBitsEighth uint,
-	tfEstimate float32,
+	tfEstimate float32, intensity int,
 ) {
 	info.allocationTrim = defaultAllocationTrim
 	if e.rangeEncoder.TellFrac()+uint(allocationTrimBitCost<<bitResolution) <= totalBitsEighth {
@@ -368,7 +372,7 @@ func (e *Encoder) encodeAllocationTrim(
 			mdct,
 			info.channelCount, info.lm, info.endBand,
 			info.totalBits,
-			tfEstimate,
+			tfEstimate, intensity, &e.stereoSaving,
 		)
 		e.rangeEncoder.EncodeSymbolWithICDF(icdfAllocationTrim, uint32(info.allocationTrim))
 	}
@@ -538,8 +542,34 @@ func computeVBR(
 	channelCount int,
 	lm int,
 	tfEstimate float32,
+	intensity, lastCodedBands int, stereoSaving float32,
 ) int {
-	target := baseTarget + totBoostBits - (19 << lm) // dynalloc calibration
+	target := baseTarget
+
+	// Stereo savings: bands coded in intensity stereo carry one channel's
+	// worth of shape, so the frame needs fewer bits the more redundant the
+	// two channels are. Capped by the share of the spectrum that is actually
+	// stereo-coded (celt_encoder.c compute_vbr).
+	if channelCount == 2 {
+		codedBands := lastCodedBands
+		if codedBands == 0 {
+			codedBands = maxBands
+		}
+		codedBins := int(bandEdges[codedBands]) << lm
+		codedStereoBands := min(intensity, codedBands)
+		codedBins += int(bandEdges[codedStereoBands]) << lm
+		codedStereoDOF := (int(bandEdges[codedStereoBands]) << lm) - codedStereoBands
+		if codedBins > 0 {
+			maxFrac := 0.8 * float32(codedStereoDOF) / float32(codedBins)
+			saving := min32(stereoSaving, 1.0)
+			target -= min(
+				int(maxFrac*float32(target)),
+				int((saving-0.1)*float32(codedStereoDOF<<bitResolution)),
+			)
+		}
+	}
+
+	target += totBoostBits - (19 << lm) // dynalloc calibration
 
 	// Transient boost, compensated for the average frame. The reference scales
 	// the target by tf_estimate rather than switching on the transient flag.
@@ -579,7 +609,7 @@ func (e *Encoder) frameCeiling(dst []byte, frameBytes int) int {
 
 func (e *Encoder) applyVBR(
 	frameBytes int, dr dynallocResult,
-	maxBytes, lm, channelCount, tellFrac int, tfEstimate float32,
+	maxBytes, lm, channelCount, tellFrac int, tfEstimate float32, intensity int,
 ) int {
 	vbrRate := frameBytes << 6 // libopus vbr_rate, 1/8-bit units
 
@@ -600,6 +630,7 @@ func (e *Encoder) applyVBR(
 	// and the rounded value for the reservoir — they're not the same number.
 	rawTarget := computeVBR(
 		baseTarget, dr.maxDepth, dr.totBoostBits, e.constrainedVBR, channelCount, lm, tfEstimate,
+		intensity, e.lastCodedBands, e.stereoSaving,
 	) + tellFrac
 
 	// The frame still has to fit what has already been written plus the
@@ -771,12 +802,17 @@ func (e *Encoder) EncodeFrame(pcm [][]float32, dst []byte, frameBytes, startBand
 	e.prevSpreadDecision = info.spread
 	e.encodeSpread(&info)
 	totalBitsEighth := e.encodeDynamicAllocation(&info, offsets)
-	e.encodeAllocationTrim(&info, analysis.logBandAmp, analysis.mdct, totalBitsEighth, tfEstimate)
+	// The reference settles the intensity band before both the trim and the VBR
+	// target, and both read it (celt_encoder.c:2404). It also derives it from
+	// the nominal rate, not from the post-VBR frame size.
+	targetIntensity, targetDualStereo := e.computeIntensityAndDualStereo(&info, normalized)
+	e.encodeAllocationTrim(&info, analysis.logBandAmp, analysis.mdct, totalBitsEighth, tfEstimate, targetIntensity)
 
 	tellFrac := int(e.rangeEncoder.TellFrac())
 
 	if e.vbr {
-		effectiveBytes = e.applyVBR(frameBytes, dr, maxBytes, info.lm, info.channelCount, tellFrac, tfEstimate)
+		effectiveBytes = e.applyVBR(
+			frameBytes, dr, maxBytes, info.lm, info.channelCount, tellFrac, tfEstimate, targetIntensity)
 	}
 	info.totalBits = uint(effectiveBytes) * 8
 	shapeBits := (int(info.totalBits) << bitResolution) - tellFrac - 1
@@ -785,7 +821,6 @@ func (e *Encoder) EncodeFrame(pcm [][]float32, dst []byte, frameBytes, startBand
 		info.antiCollapseRsv = 1 << bitResolution
 	}
 	shapeBits -= info.antiCollapseRsv
-	targetIntensity, targetDualStereo := e.computeIntensityAndDualStereo(&info, normalized)
 	info.allocation = e.computeAllocationMono(&info, shapeBits, targetIntensity, targetDualStereo)
 	e.encodeFineEnergy(&info, info.allocation.fineQuant, targetLogE)