I am getting the error "ValueError: math domain error" for the following line of the code. I am trying to take the log of some probabilities.

k0 = math.log2(probs_0[k][0])

How can I fix the error?

You are getting the error because the probability is 0. log2(0) is undefined. Make the following change to fix the error:

k0 = math.log2(probs_0[k][0] if probs_0[k][0]>0 else 1)