Problem 673306 · medium · Level 06 Heuristics & Optimization

Two Kinds of Fish

mixture model · EM algorithm · responsibilities · log-sum-exp · maximum likelihood

A lake survey measured the lengths of fish (in cm) caught in a net. Two species live in the lake, but the survey did not record which species each fish was. The biologist models a length as coming from species 1 with probability w (lengths normal with mean mu1 and standard deviation sd1) or from species 2 with probability 1 - w (mean mu2, standard deviation sd2).

Write fit_two_species(lengths, start, rounds, floor) that improves the parameters with this procedure, starting from start = (w, mu1, sd1, mu2, sd2) and repeating it exactly rounds times:

  1. For every length x, compute r(x), the probability that the fish belongs to species 1 given x and the current parameters (Bayes' rule with the two weighted normal densities).
  2. Replace the parameters: n1 is the sum of all r(x) and n2 = len(lengths) - n1; mu1 is the r-weighted mean of the lengths and sd1 the square root of the r-weighted mean of (x - mu1)**2 (with the new mu1), both dividing by n1; mu2 and sd2 likewise with weights 1 - r(x) and n2; w = n1 / len(lengths). If a new standard deviation is below floor, use floor instead.

Return a tuple (w, mu1, sd1, mu2, sd2, loglik) with the final parameters and the log-likelihood of all the lengths under them (the sum over x of log(w * f1(x) + (1 - w) * f2(x)), where f1 and f2 are the normal densities).

Some lengths may be very far from both means; your computation must still work.

The setup provides lake_catch(n, seed), which simulates n fish lengths from a lake with two species.

Examples

Input:  lengths = [1.0, 1.2, 0.8, 5.0, 5.3, 4.9, 5.1], start = (0.5, 0.0, 1.0, 6.0, 1.0),
        rounds = 1, floor = 0.01
Output: (0.42857047618696514, 1.0000275738377116, 0.16364935775059156,
         5.074972528021756, 0.14827472239385525, -1.6316325163455838)
Explanation: after a single round the three short fish are already almost
entirely assigned to species 1.

Input:  lengths = [0.0, 0.1, 1000.0], start = (0.5, 0.0, 1.0, 1.0, 1.0), rounds = 3, floor = 0.05
Output: (0.6666666666666666, 0.05, 0.05, 1000.0, 0.05, 3.3208387161635162)
Explanation: the 1000 cm reading is so far from both means that both densities
are 0.0 in floating point; the procedure must still work. Species 2 ends up with
only that fish, and its standard deviation is held at the floor.

Constraints

  • 2 <= len(lengths) <= 3000, 0 <= rounds <= 40, floor > 0, 0 < w < 1 and sd1, sd2 > 0 in start
  • the tests never let either species' total weight n1 or n2 become 0
  • floats are compared with a tolerance of 1e-6

Goals

  • Fit a two-component normal mixture by alternating soft assignments and weighted refits
  • Follow a precisely specified procedure so the result is reproducible
  • Compute responsibilities and the log-likelihood without underflow
Starting Python…