ó
ô¬VCc           @   st  d  Z  d d l Z d d l m Z d d l m Z d d l m Z d d l m Z d d l m	 Z	 d d l m
 Z
 d	 Z d
 e j f d „  ƒ  YZ d e j f d „  ƒ  YZ d „  Z d „  Z e j e ƒ  e ƒ  ƒ Z e j ƒ  e j ƒ  e j ƒ  Z e j ƒ  Z e d ƒ \ Z Z d „  Z d GHe	 j e e ƒ Z e	 j e ƒ Z e j e g ƒ Z e rle j  GHe j! GHn  e d ƒ \ Z" Z# e j$ e" e ƒ  ƒ \ Z% Z& e r¾d Ge& GHe
 j' e" e# e% ƒ n  d GHe	 j e e d e ƒ  ƒ ƒ Z( e	 j) e ƒ Z e j e( g e ƒ Z e re j  GHe j! GHn  e d ƒ \ Z" Z# e j$ e" e ƒ  ƒ \ Z% Z& e rpd Ge& GHe
 j' e" e# e% ƒ n  d S(   s	  Test out HMMs using the Occasionally Dishonest Casino.

This uses the ocassionally dishonest casino example from Biological
Sequence Analysis by Durbin et al.

In this example, we are dealing with a casino that has two types of
dice, a fair dice that has 1/6 probability of rolling any number and
a loaded dice that has 1/2 probability to roll a 6, and 1/10 probability
to roll any other number. The probability of switching from the fair to
loaded dice is .05 and the probability of switching from loaded to fair is
.1.
iÿÿÿÿN(   t   Alphabet(   t
   MutableSeq(   t   Seq(   t   MarkovModel(   t   Trainer(   t	   Utilitiesi    t   DiceRollAlphabetc           B   s    e  Z d  d d d d d g Z RS(   t   1t   2t   3t   4t   5t   6(   t   __name__t
   __module__t   letters(    (    (    s   test_HMMCasino.pyR       s   t   DiceTypeAlphabetc           B   s   e  Z d  d g Z RS(   t   Ft   L(   R   R   R   (    (    (    s   test_HMMCasino.pyR   #   s   c         C   sz  | d k r³ |  t  d ƒ t  d ƒ k r, d S|  t  d ƒ t  d ƒ k rL d S|  t  d ƒ t  d ƒ k rl d S|  t  d	 ƒ t  d ƒ k rŒ d
 S|  t  d ƒ t  d ƒ k r¬ d Sd SnÃ | d k rf|  t  d ƒ t  d ƒ k rß d S|  t  d ƒ t  d ƒ k rÿ d S|  t  d ƒ t  d ƒ k rd S|  t  d	 ƒ t  d ƒ k r?d
 S|  t  d ƒ t  d ƒ k r_d Sd Sn t d | ƒ ‚ d S(   sG   Generate a loaded dice roll based on the state and a random number
    R   i   i   R   i   R   i   R	   i   R
   i   R   R   R   i
   s   Unexpected cur_state %sN(   t   floatt
   ValueError(   t
   chance_numt	   cur_state(    (    s   test_HMMCasino.pyt   _loaded_dice_roll'   s2    c         C   sÚ   d } t  d t ƒ } t  d t ƒ } x t |  ƒ D] } | j | ƒ t j ƒ  } t | | ƒ } | j | ƒ t j ƒ  } | d k rœ | d k rÀ d } qÀ q1 | d k r1 | d k rÀ d } qÀ q1 q1 W| j ƒ  | j ƒ  f S(   s¬   Generate a bunch of rolls corresponding to the casino probabilities.

    Returns:

    o The generate roll sequence

    o The state sequence that generated the roll.
    R   t    gš™™™™™©?R   gš™™™™™¹?(   R   R   R   t   ranget   appendt   randomR   t   toseq(   t	   num_rollsR   t   roll_seqt	   state_seqt   rollR   t   new_roll(    (    s   test_HMMCasino.pyt   generate_rollsG   s     
i¸  c         C   s:   t  r d G|  GHn  |  d k  r" d S| d k r2 d Sd Sd S(   s*   Tell the training model when to stop.
    s
   ll change:g{®Gáz„?i   i
   i    N(   t   VERBOSE(   t   log_likelihood_changet   num_iterations(    (    s   test_HMMCasino.pyt   stop_training‘   s    s%   Training with the Standard Trainer...i,  s   Prediction probability:s   Training with Baum-Welch...R   (*   t   __doc__R   t   BioR    t   Bio.SeqR   R   t   Bio.HMMR   R   R   R#   R   R   R   R"   t   MarkovModelBuildert
   mm_buildert   allow_all_transitionst   set_random_probabilitiest   get_markov_modelt   baum_welch_mmt   standard_mmt   rollst   statesR&   t   TrainingSequencet   known_training_seqt   KnownStateTrainert   trainert   traint
   trained_mmt   transition_probt   emission_probt
   test_rollst   test_statest   viterbit   predicted_statest   probt   pretty_print_predictiont   training_seqt   BaumWelchTrainer(    (    (    s   test_HMMCasino.pyt   <module>   sX   	 	$

			