view yeti/em_onecolumn.yeti @ 19:f1f8c84339d0

Starting to revisit this EM logic and test in a single-column world

module em_onecolumn;

mm = load may.mathmisc;
vec = load may.vector;
mat = load may.matrix;

inRange ranges instrument note =
    note >= ranges[instrument].lowest and note <= ranges[instrument].highest;

initialise ranges templates notes =
   (instruments = keys ranges;
    {
        pitches = // z in the original. 1 per note
            vec.randoms notes,
        sources = // u in the original. 1 per note-instrument
            mapIntoHash id
                do instrument:
                    vec.fromList
                       (map do note:
                            if inRange ranges instrument note then 1 else 0 fi
                        done [0..notes-1])
                done instruments,
        instruments,
        templates,
        ranges,
        lowest = head
           (sort (map do i: ranges[i].lowest done instruments)),
        highest = head (reverse
           (sort (map do i: ranges[i].highest done instruments))),
    });

epsilon = 1e-16;

select predicate = concatMap do v: if predicate v then [v] else [] fi done;

distributionsFor data instrument note =
    {
        w = mat.getColumn note data.templates[instrument],
        p = vec.at data.pitches note,
        s = vec.at data.sources[instrument] note,
    };

performExpectation data column =
   (estimate = 
        fold do acc instrument:
            fold do acc note:
                { w, p, s } = distributionsFor data instrument note;
                vec.add [acc, vec.scaled (p * s) w];
            done acc [data.ranges[instrument].lowest .. 
                      data.ranges[instrument].highest]
        done (vec.consts epsilon (vec.length column)) data.instruments;
    vec.divide column estimate);

performMaximisation data column q =
   (pitchTop = vec.fromList
       (map do note:
            fold do acc instrument:
                { w, p, s } = distributionsFor data instrument note;
                fold do acc bin:
                    acc + s * (vec.at w bin) * (vec.at column bin);
                done acc [0..vec.length column - 1]
            done epsilon data.instruments;
        done [data.lowest .. data.highest]);
    pitchBottom = vec.sum pitchTop;

/*    sourceTop = vec.fromList
       (map do instrument:
            fold do acc note:
                { w, p, s } = distributionsFor data instrument note;
                fold do acc bin:
                    acc + p * (vec.at w bin) * (vec.at column bin);
                done acc [0..vec.length column - 1]
            done epsilon [data.lowest .. data.highest]
        done data.instruments);
    sourceBottom = vec.sum sourceTop;
*/
    data with { 
        pitches = vec.divideBy pitchBottom pitchTop,
//        sources = vec.divideBy sourceBottom sourceTop
    });

{
    initialise,
    performExpectation,
    performMaximisation,
}