Chris@13: 
Chris@13: module em;
Chris@13: 
Chris@13: mm = load may.mathmisc;
Chris@13: vec = load may.vector;
Chris@13: mat = load may.matrix;
Chris@13: 
Chris@14: inRange ranges instrument note =
Chris@14:     note >= ranges[instrument].lowest and note <= ranges[instrument].highest;
Chris@14: 
Chris@14: initialise ranges templates notes size =
Chris@14:    (instruments = keys ranges;
Chris@13:     {
Chris@14:         pitches = // z in the original. 1xN per note
Chris@13:             array (map do note:
Chris@14:                 mat.randomMatrix { rows = 1, columns = size.columns }
Chris@13:             done [0..notes-1]),
Chris@14:         sources = // u in the original. 1xN per note-instrument
Chris@14:             mapIntoHash id
Chris@13:                 do instrument:
Chris@13:                     array (map do note:
Chris@14:                         mat.constMatrix
Chris@14:                            (if inRange ranges instrument note then 1 else 0 fi)
Chris@14:                            (size with { rows = 1 })
Chris@13:                     done [0..notes-1])
Chris@14:                 done instruments,
Chris@14:         instruments,
Chris@14:         templates,
Chris@14:         ranges,
Chris@14:         lowest = head (sort (map do i: ranges[i].lowest done instruments)),
Chris@14:         highest = head (reverse (sort (map do i: ranges[i].highest done instruments))),
Chris@14:     });
Chris@14: 
Chris@14: epsilon = 1e-16;
Chris@14: 
Chris@14: select predicate = concatMap do v: if predicate v then [v] else [] fi done;
Chris@14: 
Chris@14: performExpectation data chunk =
Chris@14:    (estimate = 
Chris@14:         fold do acc instrument:
Chris@14:             fold do acc note:
Chris@14:                 template = mat.getColumn note data.templates[instrument];
Chris@15:                 resize = mat.tiledTo (mat.size chunk);
Chris@15:                 w = resize (mat.newColumnVector template);
Chris@15:                 p = resize data.pitches[note];
Chris@15:                 s = resize data.sources[instrument][note];
Chris@14:                 mat.sum [acc, mat.entryWiseProduct [w, p, s]];
Chris@14:             done acc [data.ranges[instrument].lowest .. 
Chris@14:                       data.ranges[instrument].highest]
Chris@14:         done (mat.constMatrix epsilon (mat.size chunk)) data.instruments;
Chris@14:     mat.entryWiseDivide chunk estimate);
Chris@14: 
Chris@14: performMaximisation data chunk error =
Chris@15:    (pitches =
Chris@15:         fold do acc note:
Chris@15:             fold do acc instrument:
Chris@15:                 // want sum of error * original for all template and instruments
Chris@15:                 // for this pitch, divided by sum of error * original for all
Chris@15:                 // template and instruments for all pitches
Chris@15: 
Chris@14:             template = mat.getColumn note data.templates[instrument];
Chris@14:             w = mat.repeatedHorizontal (mat.width chunk) (mat.newColumnVector template);
Chris@14:             p = mat.repeatedVertical (mat.height chunk) data.pitches[note];
Chris@14:             s = mat.repeatedVertical (mat.height chunk) data.sources[instrument][note];
Chris@15:             mat.sum [acc, mat.entryWiseProduct [w, p, s, error]]
Chris@14:         done acc (select do i: inRange data.ranges i note done data.instruments)
Chris@15:     done (mat.constMatrix epsilon  [data.lowest .. data.highest]);
Chris@13: 
Chris@13: {
Chris@14:     initialise,
Chris@14:     performExpectation,
Chris@14:     performMaximisation,
Chris@13: }
Chris@13: 
Chris@13: