% **************************************************
% Frequency Masking FIR
% Markus Nentwig, 2.10.2012
%
% based on 
% [1] http://www-ee.uta.edu/online/oraintara/ee5350/pdfs/01085930.pdf
% specifications: [1] eq. 16
% parameters: [1] table I row M=9
% topology: [1] fig. 3 and 5 
% With regard to this picture, any ..._upper label
%   in this program refers to the *lower* branch and vice versa)
% 
% This program attempts a least-squares optimal design, 
%   whereas [1] is minmax. 
% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
% **************************************************
function frMaskingFIR_121002()
    close all;

    % "object" (struct) to pass the design to the optimizer
    obj = oh_setup(struct());

    obj.refFlag = false;
    if obj.refFlag
        % **************************************************
        % "reference" filter:
        % try to achieve the same with a plain ol' FIR
        % 401 taps... that's a lot...
        % **************************************************
        obj.delay_ref = 200; % delay; length = 2*delay+1
        
        m = map_symFIR(obj.delay_ref);
        obj.map_ref = m;
        obj.ref = zeros(1, size(m, 1));
        obj = oh_setOptimVar(obj, 'ref');
        
    else
        % **************************************************
        % "Frequency response masking" design
        % "f" filter
        % **************************************************
        n_f = 22; % delay; length = 2*delay+1; "N" in [1] (2*22+1 = 45)
        os_f = 9; % delay repetition in first stage; "M" in [1]
        
        m = map_symFIR(n_f);
        m = map_upsample(m, os_f);    
        obj.map_f = m;
        obj.delay_f = n_f * os_f;    
        obj.f = zeros(1, size(m, 1));
        obj = oh_setOptimVar(obj, 'f');
        
        % **************************************************
        % "g" filters
        % same structure, different coefficients
        % **************************************************
        obj.delay_g1 = 16; % N_{MC} in [1](2*16+1 = 33)
        obj.delay_g2 = 20; % N_{MA} in [1](2*20+1 = 41)
        obj.delay_g = max(obj.delay_g1, obj.delay_g2);
        
        m = map_symFIR(obj.delay_g1);
        m = map_pad(m, obj.delay_g);
        obj.map_g1 = m;
        obj.g1 = zeros(1, size(m, 1));
        
        m = map_symFIR(obj.delay_g2);
        m = map_pad(m, obj.delay_g);
        obj.map_g2 = m;
        obj.g2 = zeros(1, size(m, 1));
        
        obj = oh_setOptimVar(obj, 'g1');
        obj = oh_setOptimVar(obj, 'g2');
    end
    
    % **************************************************
    % design problem
    % **************************************************
    % frequency grid for evaluation
    obj.nGrid = 2048;
    omega = freqbase(obj.nGrid);

    % targeted frequency response
    H = zeros(1, obj.nGrid);
    
    % passband
    mask = (abs(omega) <= 0.3);
    H(mask) = 1;
    % weighting of least-squares error
    w(mask) = 1 / (1 - 10.^(0.05 / 20));
    
    % stopband
    mask = (abs(omega) >= 0.305);
    H(mask) = 0;
    w(mask) = 1 / 10 .^ (-40/20);
    
    % known group delay of topology
    if obj.refFlag
        H = H .* exp(-1i*omega*2*pi*(obj.delay_ref));
    else
        % g1 and g2 are padded for identical group delay
        % f and g1|g2 appear in series
        H = H .* exp(-1i*omega*2*pi*(obj.delay_f + obj.delay_g));
    end

    obj.HTarget = H;
    obj.weight = w;
    
    % **************************************************
    % solve
    % **************************************************
    % wrapper to pass obj as extra parameter
    ofunWrapper = @(vec) ofun(vec, obj);        

    vec = fminunc(...
        ofunWrapper, ...
        obj.vec, ...
        optimset('tolX', 1e-6, 'maxFunEvals', 100000, 'maxIter', 2000));
    obj = oh_restoreOptimVar(obj, vec);
    
    % **************************************************
    % plot
    % **************************************************
    H = evalH(obj);
    Herr = (H-obj.HTarget);
    HerrW = Herr .* obj.weight;
    totalError_dB = s2dB(norm(HerrW))
    
    figure(1);
    hold on;
    h = plotH(omega, H, 'b'); set(h, 'lineWidth', 3);
    %plotH(omega, Herr, 'r');
    %plotH(omega, HerrW, 'r');
    %legend('|H_{combined}|', 'error', 'error (weighted)');
    
    if ~obj.refFlag
        % **************************************************
        % plot individual stages
        % **************************************************
        figure(2); hold on;
        % upper "f" stage: pure delay
        ir = [];
        ir([obj.delay_f+1, obj.nGrid]) = [1, 0];
        H1 = fft(ir);
        
        % lower "f" stage
        ir = ((obj.map_f .') * obj.f .') .';
        ir(obj.nGrid) = 0;    
        plotH(omega, fft(ir), 'k');
        plotH(omega, H1 - fft(ir), 'b');
        
        % upper "g" stage
        ir = ((obj.map_g1 .') * obj.g1 .') .';
        ir(obj.nGrid) = 0;
        plotH(omega, fft(ir), 'g');
        
        % lower half
        ir = ((obj.map_g2 .') * obj.g2 .') .';
        ir(obj.nGrid) = 0;
        plotH(omega, fft(ir), 'c');
        legend('f1', '1-f1', 'g1', 'g2');
        xlabel('f / f_{s}');
        ylabel('|H|; dB');
    end
end

% **************************************************
% objective function for optimizer
% **************************************************    
function errW = ofun(vec, obj)
% put optimizer's "vec" back into obj
    obj = oh_restoreOptimVar(obj, vec);
    
    % calculate frequency response 
    H = evalH(obj);

    % calculate per-frequency error
    err = abs(H - obj.HTarget);
    
    % weighted with requirements
    errW = norm(obj.weight .* err);
end

% **************************************************
% evaluate frequency response
% **************************************************    
function H = evalH(obj)

    if obj.refFlag
        % reference filter
      
        % construct the impulse response from the coefficient map
        % (filter topology, fixed) and the coefficient vector (varies
        % with optimization)
        ir = ((obj.map_ref .') * obj.ref .') .';

        % pad the impulse response to the FFT size / frequency grid length
        ir(obj.nGrid) = 0;

        % calculate frequency response
        H = fft(ir);
        return;
    end

    % the code below follows the above pattern
    
    % upper "f" stage: pure delay
    ir = [];
    ir([obj.delay_f+1, obj.nGrid]) = [1, 0];
    Hf_upper = fft(ir);
    
    % lower "f" stage
    ir = ((obj.map_f .') * obj.f .') .';
    ir(obj.nGrid) = 0;    
    Hf_lower = fft(ir);
    
    % upper "g" stage
    % here, the complementary "f" filter output is constructed as "1-f".
    Gin_upper = Hf_upper - Hf_lower;
    ir = ((obj.map_g1 .') * obj.g1 .') .';
    ir(obj.nGrid) = 0;
    Gout_upper = Gin_upper .* fft(ir);
    
    % lower "g" stage
    Gin_lower = Hf_lower;
    ir = ((obj.map_g2 .') * obj.g2 .') .';
    ir(obj.nGrid) = 0;
    Gout_lower = Gin_lower .* fft(ir);
    
    % output
    H = Gout_upper + Gout_lower;
end

% convert complex magnitude to dB
function s_dB = s2dB(s)
    s_dB = 20*log10(abs(s + eps));
end

function handle = plotH(omega, H, plotStyle)
    handle = plot(fftshift(omega), fftshift(s2dB(H)), plotStyle);
    set(handle, 'lineWidth', 5);
end

% calculate FFT bin frequencies
function fb = freqbase(n)
    fb = 0:(n - 1);
    fb = fb + floor(n / 2);
    fb = mod(fb, n);
    fb = fb - floor(n / 2);
    
    fb = fb / n; % now [0..0.5[, [-0.5..0[
end

% **************************************************
% /lib/map: coefficient-to-tap mapping for FIR filters
% **************************************************

% - columns correspond to delay
% - rows correspond to independent coefficients
% - in the symmetric FIR example below, the first coefficient 
%   appears in the first and last tap:
% 1 0 0 0 1
% 0 1 0 1 0
% 0 0 1 0 0
%
% the same with double delays:
% 1 0 0 0 0 0 0 0 1
% 0 0 1 0 0 0 1 0 0
% 0 0 0 0 1 0 0 0 0
%
% ordinary (asymmetric) FIR with 4 taps:
% 1 0 0 0
% 0 1 0 0
% 0 0 1 0
% 0 0 0 1

% generic FIR
function map = map_asymFIR(delay_samples)
    map = eye(2 * delay_samples + 1);
end

% symmetric FIR
function map = map_symFIR(delay_samples)
    map = zeros(delay_samples + 1, 2 * delay_samples + 1);
    src = 1 : delay_samples + 1;
    dest = 1 : delay_samples + 1;
    map(sub2ind(size(map), dest, src))= 1;
    src = delay_samples + 2 : 2 * delay_samples + 1;
    dest = delay_samples : -1 : 1;
    map(sub2ind(size(map), dest, src)) = 1;
end

% halfband FIR
function map = map_halfband(delay)
    assert(mod(delay, 2) == 1);
    map = [];
    ixc = 1;
    d = 0:2:delay-1;
    ix = 1:numel(d);
    d = [d delay d + delay + 1];
    ix = [ix, max(ix) + 1, fliplr(ix)];
    for k = 1 : numel(d)
        map(ix(k), d(k)+1) = 1;
    end
end

% interpolate a FIR map
function map = map_upsample(map, os)
    m2 = []; %zeros(size(map, 1), size(map, 2));
             
    ix1 = 1 : size(map, 2);
    ix2 = (ix1 - 1) * os + 1;
    
    m2(:, ix2) = map(:, ix1);
    map = m2;
end

% pad FIR filter with zero taps
function m = map_pad(m, d)
    sNom = 2 * d + 1;
    sAct = size(m, 2);
    nY = size(m, 1);
    nPad = (sNom - sAct) / 2;
    m = [zeros(nY, nPad) m zeros(nY, nPad)];
end

% **************************************************
% /lib/oh: "optimization helper" functions 
% bundle and separate structure entries to a flat vector
% for use with generic optimizer algorithms
% **************************************************
function obj = oh_setup(obj)
    obj.optimVarList = {};
    obj.vec = [];
end

% 'moves' variable vName into the optimization. 
% copies value to initial optim. vector vec.
% vRef: optional redundant scaling factor to bring all
% optimizer variables into the same numerical range.
% Mirror image of "restoreOptimVar"
function q = oh_setOptimVar(q, vName, vRef)
    val = q.(vName);

    if nargin < 4
        % normalize initial value to 1.
        % this won't work if it is too close to 0.
        vRef = norm(q.(vName));
    end
    if abs(vRef) < 1000*eps
        vRef = 1;
    end
    
    o = struct(...
        'name', vName, ...
        'size', size(val), ...
        'ix', numel(q.vec) + 1, ...
        'ref', vRef);
    
    q.optimVarList{end+1} = o;    
    [s1, s2] = size(val);
    
    v = val / vRef;
    q.vec = [q.vec v(:) .'];
end

% optimized variables are passed in vec, both to the 
% objective function and as optimizer result.
% Restore them into data structure q.
% Mirror image of "setOptimVar"
function q = oh_restoreOptimVar(q, vec)    
    ol = q.optimVarList;
    for ix = 1:numel(ol)
        o = ol{ix};
        ix1 = o.ix;
        ix2 = ix1 + prod(o.size) - 1;
        v = vec(ix1:ix2) .';
        v = reshape(v, o.size);
        q.(o.name) = v * o.ref;
    end
end
