bitalloc.py

"""
Music 422
-----------------------------------------------------------------------
(c) 2009-2026 Marina Bosi  -- All rights reserved
-----------------------------------------------------------------------
"""

import numpy as np
from numba import njit


# Question 1.c)
def BitAllocUniform(bitBudget, maxMantBits, nBands, nLines, SMR=None):
    """
    Returns a hard-coded vector that, in the case of the signal used in HW#4,
    gives the allocation of mantissa bits in each scale factor band when
    bits are uniformly distributed for the mantissas.
    """
    # if bitBudget > 2000:
    #     return np.array([4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4.,
    #    4., 4., 3., 4., 4., 4., 4., 4.], dtype=np.uint64)
    return np.array(
        [
            3.0,
            4.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            3.0,
            2.0,
            2.0,
        ],
        dtype=np.uint64,
    )


def BitAllocConstSNR(bitBudget, maxMantBits, nBands, nLines, peakSPL):
    """
    Returns a hard-coded vector that, in the case of the signal used in HW#4,
    gives the allocation of mantissa bits in each scale factor band when
    bits are distributed for the mantissas to try and keep a constant
    quantization noise floor (assuming a noise floor 6 dB per bit below
    the peak SPL line in the scale factor band).
    """
    # if bitBudget > 2000:
    #     return np.array([16., 16., 16., 16., 16., 15., 14., 16., 12.,  6.,  5.,  5.,  5.,
    #     4.,  4.,  4.,  5., 14., 16.,  3.,  3., 15.,  3.,  0.,  0.], dtype=np.uint64)
    return np.array(
        [
            14.0,
            16.0,
            15.0,
            16.0,
            16.0,
            12.0,
            11.0,
            15.0,
            9.0,
            3.0,
            2.0,
            2.0,
            2.0,
            3.0,
            0.0,
            0.0,
            0.0,
            11.0,
            14.0,
            0.0,
            0.0,
            12.0,
            0.0,
            0.0,
            0.0,
        ],
        dtype=np.uint64,
    )


def BitAllocConstNMR(bitBudget, maxMantBits, nBands, nLines, SMR):
    """
    Returns a hard-coded vector that, in the case of the signal used in HW#4,
    gives the allocation of mantissa bits in each scale factor band when
    bits are distributed for the mantissas to try and keep the quantization
    noise floor a constant distance below (or above, if bit starved) the
    masked threshold curve (assuming a quantization noise floor 6 dB per
    bit below the peak SPL line in the scale factor band).
    """
    #     return np.array([10., 13., 11., 12., 12., 10., 11., 14.,  7.,  5.,  7.,  9., 10.,
    #    10., 10.,  9.,  5., 11., 13.,  3.,  4., 12.,  2.,  2.,  0.], dtype=np.uint64)
    return np.array(
        [
            8.0,
            10.0,
            9.0,
            9.0,
            10.0,
            7.0,
            8.0,
            11.0,
            6.0,
            3.0,
            5.0,
            8.0,
            8.0,
            8.0,
            8.0,
            8.0,
            2.0,
            8.0,
            10.0,
            0.0,
            0.0,
            11.0,
            0.0,
            0.0,
            0.0,
        ],
        dtype=np.uint64,
    )


# Question 2.a)
@njit(cache=True)
def BitAlloc(bitBudget, maxMantBits, nBands, nLines, SMR, nmr_cutoff=-1e9):
    """
    Allocates bits to scale factor bands so as to flatten the NMR across the spectrum

       Arguments:
           bitBudget is total number of mantissa bits to allocate
           maxMantBits is max mantissa bits that can be allocated per line
           nBands is total number of scale factor bands
           nLines[nBands] is number of lines in each scale factor band
           SMR[nBands] is signal-to-mask ratio in each scale factor band

        Returns:
            bits[nBands] is number of bits allocated to each scale factor band


    """
    # return BitAllocUniform(bitBudget, maxMantBits, nBands, nLines, SMR)
    # return BitAllocConstSNR(bitBudget, maxMantBits, nBands, nLines, 0)
    # return BitAllocConstNMR(bitBudget, maxMantBits, nBands, nLines, SMR)

    # Sick water filling algorithm 9000

    bits_remaining = int(bitBudget)
    bits = np.zeros(nBands, dtype=np.uint64)
    nmr = SMR.copy()  # Initialize NMR (will update incrementally)

    while True:
        # Find band with worst (highest) NMR
        band = np.argmax(nmr)

        # Stop once worst NMR band reaches the NMR cutoff
        if nmr[band] <= nmr_cutoff:
            break

        # If band is maxed, mark it and try next band
        if bits[band] >= maxMantBits:
            nmr[band] = -1e9
            continue

        to_alloc = 2 if bits[band] == 0 else 1
        bits_used = to_alloc * nLines[band]

        if bits_remaining >= bits_used:
            bits[band] += to_alloc
            bits_remaining -= bits_used
            nmr[band] = SMR[band] - 6.0 * bits[band]  # Update only changed band
        else:
            # Mark this band as done (can't afford it)
            nmr[band] = -1e9

    return bits


# -----------------------------------------------------------------------------

# Testing code
if __name__ == "__main__":
    pass