radu
/
SSMBook
spogulis no https://github.com/probml/ssm-book.git


			
							123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979
							
<!DOCTYPE html>

<html>
  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>Hidden Markov Models &#8212; State Space Models: A Modern Approach</title>
    
  <link href="../../_static/css/theme.css" rel="stylesheet">
  <link href="../../_static/css/index.ff1ffe594081f20da1ef19478df9384b.css" rel="stylesheet">

    
  <link rel="stylesheet"
    href="../../_static/vendor/fontawesome/5.13.0/css/all.min.css">
  <link rel="preload" as="font" type="font/woff2" crossorigin
    href="../../_static/vendor/fontawesome/5.13.0/webfonts/fa-solid-900.woff2">
  <link rel="preload" as="font" type="font/woff2" crossorigin
    href="../../_static/vendor/fontawesome/5.13.0/webfonts/fa-brands-400.woff2">

    
    <link rel="stylesheet" type="text/css" href="../../_static/pygments.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/sphinx-book-theme.css?digest=c3fdc42140077d1ad13ad2f1588a4309" />
    <link rel="stylesheet" type="text/css" href="../../_static/togglebutton.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/copybutton.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/mystnb.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/sphinx-thebe.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/panels-main.c949a650a448cc0ae9fd3441c0e17fb0.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/panels-variables.06eb56fa6e07937060861dad626602ad.css" />
    
  <link rel="preload" as="script" href="../../_static/js/index.be7d3bbb2ef33a8344ce.js">

    <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
    <script src="../../_static/jquery.js"></script>
    <script src="../../_static/underscore.js"></script>
    <script src="../../_static/doctools.js"></script>
    <script src="../../_static/clipboard.min.js"></script>
    <script src="../../_static/copybutton.js"></script>
    <script>let toggleHintShow = 'Click to show';</script>
    <script>let toggleHintHide = 'Click to hide';</script>
    <script>let toggleOpenOnPrint = 'true';</script>
    <script src="../../_static/togglebutton.js"></script>
    <script>var togglebuttonSelector = '.toggle, .admonition.dropdown, .tag_hide_input div.cell_input, .tag_hide-input div.cell_input, .tag_hide_output div.cell_output, .tag_hide-output div.cell_output, .tag_hide_cell.cell, .tag_hide-cell.cell';</script>
    <script src="../../_static/sphinx-book-theme.d59cb220de22ca1c485ebbdc042f0030.js"></script>
    <script>const THEBE_JS_URL = "https://unpkg.com/thebe@0.8.2/lib/index.js"
const thebe_selector = ".thebe,.cell"
const thebe_selector_input = "pre"
const thebe_selector_output = ".output, .cell_output"
</script>
    <script async="async" src="../../_static/sphinx-thebe.js"></script>
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" />
    <meta name="viewport" content="width=device-width, initial-scale=1" />
    <meta name="docsearch:language" content="None">
    

    <!-- Google Analytics -->
    
  </head>
  <body data-spy="scroll" data-target="#bd-toc-nav" data-offset="80">
    
    <div class="container-fluid" id="banner"></div>

    
    <div class="container-xl">
      <div class="row">
          
<div class="col-12 col-md-3 bd-sidebar site-navigation show" id="site-navigation">
    
        <div class="navbar-brand-box">
    <a class="navbar-brand text-wrap" href="../../index.html">
      
      
      <h1 class="site-logo" id="site-title">State Space Models: A Modern Approach</h1>
      
    </a>
</div><form class="bd-search d-flex align-items-center" action="../../search.html" method="get">
  <i class="icon fas fa-search"></i>
  <input type="search" class="form-control" name="q" id="search-input" placeholder="Search this book..." aria-label="Search this book..." autocomplete="off" >
</form><nav class="bd-links" id="bd-docs-nav" aria-label="Main">
    <div class="bd-toc-item active">
        <ul class="nav bd-sidenav">
 <li class="toctree-l1">
  <a class="reference internal" href="../../root.html">
   State Space Models: A Modern Approach
  </a>
 </li>
</ul>
<ul class="nav bd-sidenav">
 <li class="toctree-l1">
  <a class="reference internal" href="../scratch.html">
   Scratchpad
  </a>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../ssm/ssm_index.html">
   State Space Models
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" type="checkbox"/>
  <label for="toctree-checkbox-1">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/ssm_intro.html">
     What are State Space Models?
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/hmm.html">
     Hidden Markov Models
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/lds.html">
     Linear Gaussian SSMs
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/nlds.html">
     Nonlinear Gaussian SSMs
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/inference.html">
     Inferential goals
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="hmm_index.html">
   Hidden Markov Models
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" type="checkbox"/>
  <label for="toctree-checkbox-2">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_filter.html">
     HMM filtering (forwards algorithm)
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_smoother.html">
     HMM smoothing (forwards-backwards algorithm)
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_viterbi.html">
     Viterbi algorithm
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_parallel.html">
     Parallel HMM  smoothing
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_sampling.html">
     Forwards-filtering backwards-sampling algorithm
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../lgssm/lgssm_index.html">
   Linear-Gaussian SSMs
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" type="checkbox"/>
  <label for="toctree-checkbox-3">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_filter.html">
     Kalman filtering
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_smoother.html">
     Kalman (RTS) smoother
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_parallel.html">
     Parallel Kalman Smoother
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_sampling.html">
     Forwards-filtering backwards sampling
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../extended/extended_index.html">
   Extended (linearized) methods
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" type="checkbox"/>
  <label for="toctree-checkbox-4">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../extended/extended_filter.html">
     Extended Kalman filtering
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../extended/extended_smoother.html">
     Extended Kalman smoother
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../extended/extended_parallel.html">
     Parallel extended Kalman smoothing
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../unscented/unscented_index.html">
   Unscented methods
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" type="checkbox"/>
  <label for="toctree-checkbox-5">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../unscented/unscented_filter.html">
     Unscented filtering
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../unscented/unscented_smoother.html">
     Unscented smoothing
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../quadrature/quadrature_index.html">
   Quadrature and cubature methods
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../postlin/postlin_index.html">
   Posterior linearization
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../adf/adf_index.html">
   Assumed Density Filtering
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../vi/vi_index.html">
   Variational inference
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../pf/pf_index.html">
   Particle filtering
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../smc/smc_index.html">
   Sequential Monte Carlo
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../learning/learning_index.html">
   Offline parameter estimation (learning)
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../tracking/tracking_index.html">
   Multi-target tracking
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../ensemble/ensemble_index.html">
   Data assimilation using Ensemble Kalman filter
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../bnp/bnp_index.html">
   Bayesian non-parametric SSMs
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../changepoint/changepoint_index.html">
   Changepoint detection
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../timeseries/timeseries_index.html">
   Timeseries forecasting
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../gp/gp_index.html">
   Markovian Gaussian processes
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../ode/ode_index.html">
   Differential equations and SSMs
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../control/control_index.html">
   Optimal control
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../../bib.html">
   Bibliography
  </a>
 </li>
</ul>

    </div>
</nav> <!-- To handle the deprecated key -->

<div class="navbar_extra_footer">
  Powered by <a href="https://jupyterbook.org">Jupyter Book</a>
</div>

</div>


<main class="col py-md-3 pl-md-4 bd-content overflow-auto" role="main">
    
    <div class="topbar container-xl fixed-top">
    <div class="topbar-contents row">
        <div class="col-12 col-md-3 bd-topbar-whitespace site-navigation show"></div>
        <div class="col pl-md-4 topbar-main">
            
            <button id="navbar-toggler" class="navbar-toggler ml-0" type="button" data-toggle="collapse"
                data-toggle="tooltip" data-placement="bottom" data-target=".site-navigation" aria-controls="navbar-menu"
                aria-expanded="true" aria-label="Toggle navigation" aria-controls="site-navigation"
                title="Toggle navigation" data-toggle="tooltip" data-placement="left">
                <i class="fas fa-bars"></i>
                <i class="fas fa-arrow-left"></i>
                <i class="fas fa-arrow-up"></i>
            </button>
            
            
<div class="dropdown-buttons-trigger">
    <button id="dropdown-buttons-trigger" class="btn btn-secondary topbarbtn" aria-label="Download this page"><i
            class="fas fa-download"></i></button>

    <div class="dropdown-buttons">
        <!-- ipynb file if we had a myst markdown file -->
        
        <!-- Download raw file -->
        <a class="dropdown-buttons" href="../../_sources/chapters/hmm/hmm.ipynb"><button type="button"
                class="btn btn-secondary topbarbtn" title="Download source file" data-toggle="tooltip"
                data-placement="left">.ipynb</button></a>
        <!-- Download PDF via print -->
        <button type="button" id="download-print" class="btn btn-secondary topbarbtn" title="Print to PDF"
                onclick="printPdf(this)" data-toggle="tooltip" data-placement="left">.pdf</button>
    </div>
</div>

            <!-- Source interaction buttons -->

<div class="dropdown-buttons-trigger">
    <button id="dropdown-buttons-trigger" class="btn btn-secondary topbarbtn"
        aria-label="Connect with source repository"><i class="fab fa-github"></i></button>
    <div class="dropdown-buttons sourcebuttons">
        <a class="repository-button"
            href="https://github.com/probml/ssm-book"><button type="button" class="btn btn-secondary topbarbtn"
                data-toggle="tooltip" data-placement="left" title="Source repository"><i
                    class="fab fa-github"></i>repository</button></a>
        <a class="issues-button"
            href="https://github.com/probml/ssm-book/issues/new?title=Issue%20on%20page%20%2Fchapters/hmm/hmm.html&body=Your%20issue%20content%20here."><button
                type="button" class="btn btn-secondary topbarbtn" data-toggle="tooltip" data-placement="left"
                title="Open an issue"><i class="fas fa-lightbulb"></i>open issue</button></a>
        
    </div>
</div>

            <!-- Full screen (wrap in <a> to have style consistency -->

<a class="full-screen-button"><button type="button" class="btn btn-secondary topbarbtn" data-toggle="tooltip"
        data-placement="bottom" onclick="toggleFullScreen()" aria-label="Fullscreen mode"
        title="Fullscreen mode"><i
            class="fas fa-expand"></i></button></a>

            <!-- Launch buttons -->

<div class="dropdown-buttons-trigger">
    <button id="dropdown-buttons-trigger" class="btn btn-secondary topbarbtn"
        aria-label="Launch interactive content"><i class="fas fa-rocket"></i></button>
    <div class="dropdown-buttons">
        
        <a class="binder-button" href="https://mybinder.org/v2/gh/probml/ssm-book/main?urlpath=tree/chapters/hmm/hmm.ipynb"><button type="button"
                class="btn btn-secondary topbarbtn" title="Launch Binder" data-toggle="tooltip"
                data-placement="left"><img class="binder-button-logo"
                    src="../../_static/images/logo_binder.svg"
                    alt="Interact on binder">Binder</button></a>
        
        
        <a class="colab-button" href="https://colab.research.google.com/github/probml/ssm-book/blob/main/chapters/hmm/hmm.ipynb"><button type="button" class="btn btn-secondary topbarbtn"
                title="Launch Colab" data-toggle="tooltip" data-placement="left"><img class="colab-button-logo"
                    src="../../_static/images/logo_colab.png"
                    alt="Interact on Colab">Colab</button></a>
        
        
    </div>
</div>

        </div>

        <!-- Table of contents -->
        <div class="d-none d-md-block col-md-2 bd-toc show noprint">
            
            <div class="tocsection onthispage pt-5 pb-3">
                <i class="fas fa-list"></i> Contents
            </div>
            <nav id="bd-toc-nav" aria-label="Page">
                <ul class="visible nav section-nav flex-column">
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#boilerplate">
   Boilerplate
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#utility-code">
   Utility code
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#example-casino-hmm">
   Example: Casino HMM
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#sampling-from-the-joint">
   Sampling from the joint
  </a>
  <ul class="nav section-nav flex-column">
   <li class="toc-h3 nav-item toc-entry">
    <a class="reference internal nav-link" href="#numpy-version">
     Numpy version
    </a>
   </li>
   <li class="toc-h3 nav-item toc-entry">
    <a class="reference internal nav-link" href="#jax-version">
     JAX version
    </a>
   </li>
   <li class="toc-h3 nav-item toc-entry">
    <a class="reference internal nav-link" href="#check-correctness-by-computing-empirical-pairwise-statistics">
     Check correctness by computing empirical pairwise statistics
    </a>
   </li>
  </ul>
 </li>
</ul>

            </nav>
        </div>
    </div>
</div>
    <div id="main-content" class="row">
        <div class="col-12 col-md-9 pl-md-3 pr-md-0">
            <!-- Table of contents that is only displayed when printing the page -->
            <div id="jb-print-docs-body" class="onlyprint">
                <h1>Hidden Markov Models</h1>
                <!-- Table of contents -->
                <div id="print-main-content">
                    <div id="jb-print-toc">
                        
                        <div>
                            <h2> Contents </h2>
                        </div>
                        <nav aria-label="Page">
                            <ul class="visible nav section-nav flex-column">
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#boilerplate">
   Boilerplate
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#utility-code">
   Utility code
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#example-casino-hmm">
   Example: Casino HMM
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#sampling-from-the-joint">
   Sampling from the joint
  </a>
  <ul class="nav section-nav flex-column">
   <li class="toc-h3 nav-item toc-entry">
    <a class="reference internal nav-link" href="#numpy-version">
     Numpy version
    </a>
   </li>
   <li class="toc-h3 nav-item toc-entry">
    <a class="reference internal nav-link" href="#jax-version">
     JAX version
    </a>
   </li>
   <li class="toc-h3 nav-item toc-entry">
    <a class="reference internal nav-link" href="#check-correctness-by-computing-empirical-pairwise-statistics">
     Check correctness by computing empirical pairwise statistics
    </a>
   </li>
  </ul>
 </li>
</ul>

                        </nav>
                    </div>
                </div>
            </div>
            
              <div>
                
  <div class="tex2jax_ignore mathjax_ignore section" id="hidden-markov-models">
<span id="sec-hmm-ex"></span><h1>Hidden Markov Models<a class="headerlink" href="#hidden-markov-models" title="Permalink to this headline">¶</a></h1>
<p>In this section, we introduce Hidden Markov Models (HMMs).</p>
<div class="section" id="boilerplate">
<h2>Boilerplate<a class="headerlink" href="#boilerplate" title="Permalink to this headline">¶</a></h2>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># Install necessary libraries</span>

<span class="k">try</span><span class="p">:</span>
    <span class="kn">import</span> <span class="nn">jax</span>
<span class="k">except</span><span class="p">:</span>
    <span class="c1"># For cuda version, see https://github.com/google/jax#installation</span>
    <span class="o">%</span><span class="k">pip</span> install --upgrade &quot;jax[cpu]&quot; 
    <span class="kn">import</span> <span class="nn">jax</span>

<span class="k">try</span><span class="p">:</span>
    <span class="kn">import</span> <span class="nn">jsl</span>
<span class="k">except</span><span class="p">:</span>
    <span class="o">%</span><span class="k">pip</span> install git+https://github.com/probml/jsl
    <span class="kn">import</span> <span class="nn">jsl</span>

<span class="k">try</span><span class="p">:</span>
    <span class="kn">import</span> <span class="nn">rich</span>
<span class="k">except</span><span class="p">:</span>
    <span class="o">%</span><span class="k">pip</span> install rich
    <span class="kn">import</span> <span class="nn">rich</span>
</pre></div>
</div>
</div>
</div>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># Import standard libraries</span>

<span class="kn">import</span> <span class="nn">abc</span>
<span class="kn">from</span> <span class="nn">dataclasses</span> <span class="kn">import</span> <span class="n">dataclass</span>
<span class="kn">import</span> <span class="nn">functools</span>
<span class="kn">import</span> <span class="nn">itertools</span>

<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">NamedTuple</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span>

<span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>


<span class="kn">import</span> <span class="nn">jax</span>
<span class="kn">import</span> <span class="nn">jax.numpy</span> <span class="k">as</span> <span class="nn">jnp</span>
<span class="kn">from</span> <span class="nn">jax</span> <span class="kn">import</span> <span class="n">lax</span><span class="p">,</span> <span class="n">vmap</span><span class="p">,</span> <span class="n">jit</span><span class="p">,</span> <span class="n">grad</span>
<span class="kn">from</span> <span class="nn">jax.scipy.special</span> <span class="kn">import</span> <span class="n">logit</span>
<span class="kn">from</span> <span class="nn">jax.nn</span> <span class="kn">import</span> <span class="n">softmax</span>
<span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">partial</span>
<span class="kn">from</span> <span class="nn">jax.random</span> <span class="kn">import</span> <span class="n">PRNGKey</span><span class="p">,</span> <span class="n">split</span>

<span class="kn">import</span> <span class="nn">inspect</span>
<span class="kn">import</span> <span class="nn">inspect</span> <span class="k">as</span> <span class="nn">py_inspect</span>
<span class="kn">from</span> <span class="nn">rich</span> <span class="kn">import</span> <span class="n">inspect</span> <span class="k">as</span> <span class="n">r_inspect</span>
<span class="kn">from</span> <span class="nn">rich</span> <span class="kn">import</span> <span class="nb">print</span> <span class="k">as</span> <span class="n">r_print</span>

<span class="k">def</span> <span class="nf">print_source</span><span class="p">(</span><span class="n">fname</span><span class="p">):</span>
    <span class="n">r_print</span><span class="p">(</span><span class="n">py_inspect</span><span class="o">.</span><span class="n">getsource</span><span class="p">(</span><span class="n">fname</span><span class="p">))</span>
</pre></div>
</div>
</div>
</div>
</div>
<div class="section" id="utility-code">
<h2>Utility code<a class="headerlink" href="#utility-code" title="Permalink to this headline">¶</a></h2>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="mf">1e-15</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Normalizes the values within the axis in a way that they sum up to 1.</span>
<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    u : array</span>
<span class="sd">    axis : int</span>
<span class="sd">    eps : float</span>
<span class="sd">        Threshold for the alpha values</span>
<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    * array</span>
<span class="sd">        Normalized version of the given matrix</span>
<span class="sd">    * array(seq_len, n_hidden) :</span>
<span class="sd">        The values of the normalizer</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="n">u</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">u</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">u</span> <span class="o">&lt;</span> <span class="n">eps</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">u</span><span class="p">))</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">u</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="n">axis</span><span class="p">)</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">c</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">c</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">u</span> <span class="o">/</span> <span class="n">c</span><span class="p">,</span> <span class="n">c</span>
</pre></div>
</div>
</div>
</div>
</div>
<div class="section" id="example-casino-hmm">
<span id="sec-casino-ex"></span><h2>Example: Casino HMM<a class="headerlink" href="#example-casino-hmm" title="Permalink to this headline">¶</a></h2>
<p>We first create the “Ocassionally dishonest casino” model from <span id="id1">[<a class="reference internal" href="../../bib.html#id3" title="R. Durbin, S. Eddy, A. Krogh, and G. Mitchison. Biological Sequence Analysis: Probabilistic Models of Proteins and Nucleic Acids. Cambridge University Press, 1998.">DEKM98</a>]</span>.</p>
<div class="figure align-default" id="casino-fig">
<a class="reference internal image-reference" href="../../_images/casino.png"><img alt="../../_images/casino.png" src="../../_images/casino.png" style="width: 208.5px; height: 142.5px;" /></a>
<p class="caption"><span class="caption-text">Illustration of the casino HMM.</span><a class="headerlink" href="#casino-fig" title="Permalink to this image">¶</a></p>
</div>
<p>There are 2 hidden states, each of which emit 6 possible observations.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># state transition matrix</span>
<span class="n">A</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span>
    <span class="p">[</span><span class="mf">0.95</span><span class="p">,</span> <span class="mf">0.05</span><span class="p">],</span>
    <span class="p">[</span><span class="mf">0.10</span><span class="p">,</span> <span class="mf">0.90</span><span class="p">]</span>
<span class="p">])</span>

<span class="c1"># observation matrix</span>
<span class="n">B</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span>
    <span class="p">[</span><span class="mi">1</span><span class="o">/</span><span class="mi">6</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">6</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">6</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">6</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">6</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">6</span><span class="p">],</span> <span class="c1"># fair die</span>
    <span class="p">[</span><span class="mi">1</span><span class="o">/</span><span class="mi">10</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">10</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">10</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">10</span><span class="p">,</span> <span class="mi">1</span><span class="o">/</span><span class="mi">10</span><span class="p">,</span> <span class="mi">5</span><span class="o">/</span><span class="mi">10</span><span class="p">]</span> <span class="c1"># loaded die</span>
<span class="p">])</span>

<span class="n">pi</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">normalize</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">]))</span>
<span class="n">pi</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">pi</span><span class="p">)</span>


<span class="p">(</span><span class="n">nstates</span><span class="p">,</span> <span class="n">nobs</span><span class="p">)</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">B</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="cell_output docutils container">
<div class="output stderr highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
</pre></div>
</div>
</div>
</div>
<p>Let’s make a little data structure to store all the parameters.
We use NamedTuple rather than dataclass, since we assume these are immutable.
(Also, standard python dataclass does not work well with JAX, which requires parameters to be
pytrees, as discussed in <a class="reference external" href="https://github.com/google/jax/issues/2371">https://github.com/google/jax/issues/2371</a>).</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">Array</span> <span class="o">=</span> <span class="n">Union</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">]</span>

<span class="k">class</span> <span class="nc">HMM</span><span class="p">(</span><span class="n">NamedTuple</span><span class="p">):</span>
    <span class="n">trans_mat</span><span class="p">:</span> <span class="n">Array</span>  <span class="c1"># A : (n_states, n_states)</span>
    <span class="n">obs_mat</span><span class="p">:</span> <span class="n">Array</span>  <span class="c1"># B : (n_states, n_obs)</span>
    <span class="n">init_dist</span><span class="p">:</span> <span class="n">Array</span>  <span class="c1"># pi : (n_states)</span>

<span class="n">params_np</span> <span class="o">=</span> <span class="n">HMM</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">pi</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">params_np</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">params_np</span><span class="o">.</span><span class="n">trans_mat</span><span class="p">))</span>


<span class="n">params</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">tree_map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">params_np</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">params</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">params</span><span class="o">.</span><span class="n">trans_mat</span><span class="p">))</span>
</pre></div>
</div>
</div>
<div class="cell_output docutils container">
<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>HMM(trans_mat=array([[0.95, 0.05],
       [0.1 , 0.9 ]]), obs_mat=array([[0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667,
        0.16666667],
       [0.1       , 0.1       , 0.1       , 0.1       , 0.1       ,
        0.5       ]]), init_dist=array([0.5, 0.5], dtype=float32))
&lt;class &#39;numpy.ndarray&#39;&gt;
HMM(trans_mat=DeviceArray([[0.95, 0.05],
             [0.1 , 0.9 ]], dtype=float32), obs_mat=DeviceArray([[0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667,
              0.16666667],
             [0.1       , 0.1       , 0.1       , 0.1       , 0.1       ,
              0.5       ]], dtype=float32), init_dist=DeviceArray([0.5, 0.5], dtype=float32))
&lt;class &#39;jaxlib.xla_extension.DeviceArray&#39;&gt;
</pre></div>
</div>
</div>
</div>
</div>
<div class="section" id="sampling-from-the-joint">
<h2>Sampling from the joint<a class="headerlink" href="#sampling-from-the-joint" title="Permalink to this headline">¶</a></h2>
<p>Let’s write code to sample from this model.</p>
<div class="section" id="numpy-version">
<h3>Numpy version<a class="headerlink" href="#numpy-version" title="Permalink to this headline">¶</a></h3>
<p>First we code it in numpy using a for loop.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">hmm_sample_np</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">random_state</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="n">random_state</span><span class="p">)</span>
    <span class="n">trans_mat</span><span class="p">,</span> <span class="n">obs_mat</span><span class="p">,</span> <span class="n">init_dist</span> <span class="o">=</span> <span class="n">params</span><span class="o">.</span><span class="n">trans_mat</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">obs_mat</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">init_dist</span>
    <span class="n">n_states</span><span class="p">,</span> <span class="n">n_obs</span> <span class="o">=</span> <span class="n">obs_mat</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">state_seq</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="nb">int</span><span class="p">)</span>
    <span class="n">obs_seq</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="nb">int</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">seq_len</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">t</span><span class="o">==</span><span class="mi">0</span><span class="p">:</span>
            <span class="n">zt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">n_states</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">init_dist</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">zt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">n_states</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">trans_mat</span><span class="p">[</span><span class="n">zt</span><span class="p">])</span>
        <span class="n">yt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">n_obs</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">obs_mat</span><span class="p">[</span><span class="n">zt</span><span class="p">])</span>
        <span class="n">state_seq</span><span class="p">[</span><span class="n">t</span><span class="p">]</span> <span class="o">=</span> <span class="n">zt</span>
        <span class="n">obs_seq</span><span class="p">[</span><span class="n">t</span><span class="p">]</span> <span class="o">=</span> <span class="n">yt</span>

    <span class="k">return</span> <span class="n">state_seq</span><span class="p">,</span> <span class="n">obs_seq</span>
</pre></div>
</div>
</div>
</div>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">seq_len</span> <span class="o">=</span> <span class="mi">100</span>
<span class="n">state_seq</span><span class="p">,</span> <span class="n">obs_seq</span> <span class="o">=</span> <span class="n">hmm_sample_np</span><span class="p">(</span><span class="n">params_np</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">random_state</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">state_seq</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">obs_seq</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="cell_output docutils container">
<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 0 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[4 1 0 2 3 4 5 4 3 1 5 4 5 0 5 2 5 3 5 4 5 5 4 2 1 4 1 0 0 4 2 2 3 3 3 0 4
 0 2 4 3 2 5 5 3 5 3 1 3 3 3 2 3 5 5 0 4 4 5 0 0 1 3 5 1 5 0 1 2 4 0 0 0 4
 0 5 1 4 3 5 4 5 0 2 3 5 2 4 1 2 1 0 4 3 5 0 4 5 1 5]
</pre></div>
</div>
</div>
</div>
</div>
<div class="section" id="jax-version">
<h3>JAX version<a class="headerlink" href="#jax-version" title="Permalink to this headline">¶</a></h3>
<p>Now let’s write a JAX version using jax.lax.scan (for the inter-dependent states) and vmap (for the observations).
This is harder to read than the numpy version, but faster.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1">#@partial(jit, static_argnums=(1,))</span>
<span class="k">def</span> <span class="nf">markov_chain_sample</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">init_dist</span><span class="p">,</span> <span class="n">trans_mat</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">):</span>
    <span class="n">n_states</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">init_dist</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">draw_state</span><span class="p">(</span><span class="n">prev_state</span><span class="p">,</span> <span class="n">key</span><span class="p">):</span>
        <span class="n">state</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">n_states</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">trans_mat</span><span class="p">[</span><span class="n">prev_state</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">state</span><span class="p">,</span> <span class="n">state</span>

    <span class="n">rng_key</span><span class="p">,</span> <span class="n">rng_state</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
    <span class="n">keys</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_state</span><span class="p">,</span> <span class="n">seq_len</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">initial_state</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">n_states</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">init_dist</span><span class="p">)</span>
    <span class="n">final_state</span><span class="p">,</span> <span class="n">states</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">lax</span><span class="o">.</span><span class="n">scan</span><span class="p">(</span><span class="n">draw_state</span><span class="p">,</span> <span class="n">initial_state</span><span class="p">,</span> <span class="n">keys</span><span class="p">)</span>
    <span class="n">state_seq</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">initial_state</span><span class="p">]),</span> <span class="n">states</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">state_seq</span>
</pre></div>
</div>
</div>
</div>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1">#@partial(jit, static_argnums=(1,))</span>
<span class="k">def</span> <span class="nf">hmm_sample</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">):</span>

    <span class="n">trans_mat</span><span class="p">,</span> <span class="n">obs_mat</span><span class="p">,</span> <span class="n">init_dist</span> <span class="o">=</span> <span class="n">params</span><span class="o">.</span><span class="n">trans_mat</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">obs_mat</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">init_dist</span>
    <span class="n">n_states</span><span class="p">,</span> <span class="n">n_obs</span> <span class="o">=</span> <span class="n">obs_mat</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">rng_key</span><span class="p">,</span> <span class="n">rng_obs</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
    <span class="n">state_seq</span> <span class="o">=</span> <span class="n">markov_chain_sample</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">init_dist</span><span class="p">,</span> <span class="n">trans_mat</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">draw_obs</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="n">key</span><span class="p">):</span>
        <span class="n">obs</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">n_obs</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">obs_mat</span><span class="p">[</span><span class="n">z</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">obs</span>

    <span class="n">keys</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_obs</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
    <span class="n">obs_seq</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">draw_obs</span><span class="p">,</span> <span class="n">in_axes</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">))(</span><span class="n">state_seq</span><span class="p">,</span> <span class="n">keys</span><span class="p">)</span>
    
    <span class="k">return</span> <span class="n">state_seq</span><span class="p">,</span> <span class="n">obs_seq</span>
</pre></div>
</div>
</div>
</div>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1">#@partial(jit, static_argnums=(1,))</span>
<span class="k">def</span> <span class="nf">hmm_sample2</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">):</span>

    <span class="n">trans_mat</span><span class="p">,</span> <span class="n">obs_mat</span><span class="p">,</span> <span class="n">init_dist</span> <span class="o">=</span> <span class="n">params</span><span class="o">.</span><span class="n">trans_mat</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">obs_mat</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">init_dist</span>
    <span class="n">n_states</span><span class="p">,</span> <span class="n">n_obs</span> <span class="o">=</span> <span class="n">obs_mat</span><span class="o">.</span><span class="n">shape</span>

    <span class="k">def</span> <span class="nf">draw_state</span><span class="p">(</span><span class="n">prev_state</span><span class="p">,</span> <span class="n">key</span><span class="p">):</span>
        <span class="n">state</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">n_states</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">trans_mat</span><span class="p">[</span><span class="n">prev_state</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">state</span><span class="p">,</span> <span class="n">state</span>

    <span class="n">rng_key</span><span class="p">,</span> <span class="n">rng_state</span><span class="p">,</span> <span class="n">rng_obs</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
    <span class="n">keys</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_state</span><span class="p">,</span> <span class="n">seq_len</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">initial_state</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">n_states</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">init_dist</span><span class="p">)</span>
    <span class="n">final_state</span><span class="p">,</span> <span class="n">states</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">lax</span><span class="o">.</span><span class="n">scan</span><span class="p">(</span><span class="n">draw_state</span><span class="p">,</span> <span class="n">initial_state</span><span class="p">,</span> <span class="n">keys</span><span class="p">)</span>
    <span class="n">state_seq</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">initial_state</span><span class="p">]),</span> <span class="n">states</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">draw_obs</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="n">key</span><span class="p">):</span>
        <span class="n">obs</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">n_obs</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">obs_mat</span><span class="p">[</span><span class="n">z</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">obs</span>

    <span class="n">keys</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_obs</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
    <span class="n">obs_seq</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">draw_obs</span><span class="p">,</span> <span class="n">in_axes</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">))(</span><span class="n">state_seq</span><span class="p">,</span> <span class="n">keys</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">state_seq</span><span class="p">,</span> <span class="n">obs_seq</span>
</pre></div>
</div>
</div>
</div>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">key</span> <span class="o">=</span> <span class="n">PRNGKey</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">100</span>

<span class="n">state_seq</span><span class="p">,</span> <span class="n">obs_seq</span> <span class="o">=</span> <span class="n">hmm_sample</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">state_seq</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">obs_seq</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="cell_output docutils container">
<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>[1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1
 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[5 5 2 2 0 0 0 1 3 3 2 2 5 1 5 1 0 2 2 4 2 5 1 5 5 0 0 4 2 4 3 2 3 4 1 0 5
 2 2 2 1 4 3 2 2 2 4 1 0 3 5 2 5 1 4 2 5 2 5 0 5 4 4 4 2 2 0 4 5 2 2 0 1 5
 1 3 4 5 1 5 0 5 1 5 1 2 4 5 3 4 5 4 0 4 0 2 4 5 3 3]
</pre></div>
</div>
</div>
</div>
</div>
<div class="section" id="check-correctness-by-computing-empirical-pairwise-statistics">
<h3>Check correctness by computing empirical pairwise statistics<a class="headerlink" href="#check-correctness-by-computing-empirical-pairwise-statistics" title="Permalink to this headline">¶</a></h3>
<p>We will compute the number of i-&gt;j transitions, and check that it is close to the true
A[i,j] transition probabilites.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">collections</span>
<span class="k">def</span> <span class="nf">compute_counts</span><span class="p">(</span><span class="n">state_seq</span><span class="p">,</span> <span class="n">nstates</span><span class="p">):</span>
    <span class="n">wseq</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">state_seq</span><span class="p">)</span>
    <span class="n">word_pairs</span> <span class="o">=</span> <span class="p">[</span><span class="n">pair</span> <span class="k">for</span> <span class="n">pair</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">wseq</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">wseq</span><span class="p">[</span><span class="mi">1</span><span class="p">:])]</span>
    <span class="n">counter_pairs</span> <span class="o">=</span> <span class="n">collections</span><span class="o">.</span><span class="n">Counter</span><span class="p">(</span><span class="n">word_pairs</span><span class="p">)</span>
    <span class="n">counts</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">nstates</span><span class="p">,</span> <span class="n">nstates</span><span class="p">))</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">k</span><span class="p">,</span><span class="n">v</span><span class="p">)</span> <span class="ow">in</span> <span class="n">counter_pairs</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="n">counts</span><span class="p">[</span><span class="n">k</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">k</span><span class="p">[</span><span class="mi">1</span><span class="p">]]</span> <span class="o">=</span> <span class="n">v</span>
    <span class="k">return</span> <span class="n">counts</span>

<span class="k">def</span> <span class="nf">normalize_counts</span><span class="p">(</span><span class="n">counts</span><span class="p">):</span>
    <span class="n">ncounts</span> <span class="o">=</span> <span class="n">vmap</span><span class="p">(</span><span class="k">lambda</span> <span class="n">v</span><span class="p">:</span> <span class="n">normalize</span><span class="p">(</span><span class="n">v</span><span class="p">)[</span><span class="mi">0</span><span class="p">],</span> <span class="n">in_axes</span><span class="o">=</span><span class="mi">0</span><span class="p">)(</span><span class="n">counts</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ncounts</span>

<span class="n">init_dist</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">])</span>
<span class="n">trans_mat</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mf">0.7</span><span class="p">,</span> <span class="mf">0.3</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">]])</span>
<span class="n">rng_key</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">PRNGKey</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">500</span>
<span class="n">state_seq</span> <span class="o">=</span> <span class="n">markov_chain_sample</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">init_dist</span><span class="p">,</span> <span class="n">trans_mat</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">state_seq</span><span class="p">)</span>

<span class="n">counts</span> <span class="o">=</span> <span class="n">compute_counts</span><span class="p">(</span><span class="n">state_seq</span><span class="p">,</span> <span class="n">nstates</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">counts</span><span class="p">)</span>

<span class="n">trans_mat_empirical</span> <span class="o">=</span> <span class="n">normalize_counts</span><span class="p">(</span><span class="n">counts</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">trans_mat_empirical</span><span class="p">)</span>

<span class="k">assert</span> <span class="n">jnp</span><span class="o">.</span><span class="n">allclose</span><span class="p">(</span><span class="n">trans_mat</span><span class="p">,</span> <span class="n">trans_mat_empirical</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="mf">1e-1</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="cell_output docutils container">
<div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>[0 0 1 1 1 1 0 0 1 1 1 0 1 0 0 1 1 1 1 0 1 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 1
 1 0 0 0 0 1 0 1 0 0 0 0 1 0 0 1 1 0 1 1 0 1 1 0 1 1 1 0 0 1 1 0 1 0 0 1 0
 1 0 0 0 0 0 0 0 0 0 0 1 1 1 1 0 0 0 1 1 1 0 1 0 0 0 0 1 0 0 0 0 1 1 0 0 0
 0 0 1 1 1 1 1 1 0 0 0 1 1 0 0 0 0 0 1 0 0 0 1 0 1 1 0 1 1 0 0 0 0 0 0 1 0
 0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 1 1 1 1 0 1 1 0 0 0 0 0 1 0 0 0 0
 0 0 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 1 0 0 0 1 1 1 0 0 0 1 1 0 0 0
 0 0 0 1 1 1 0 0 0 0 1 0 0 1 1 1 0 1 1 1 1 1 0 1 1 0 0 0 1 1 0 1 0 0 1 0 0
 0 0 0 1 0 0 0 1 0 1 0 0 0 0 1 0 0 1 0 0 0 1 1 0 0 0 0 0 0 0 1 0 0 1 1 1 1
 1 1 0 0 0 0 0 1 1 0 0 0 0 0 0 1 0 0 1 0 0 0 0 0 0 1 0 0 0 0 1 0 1 0 0 0 1
 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 0 0 1 0 0 1 1 0 1 0 0 0
 0 0 0 0 0 0 0 1 0 0 1 1 1 1 0 0 1 1 0 0 0 0 1 1 0 1 1 0 0 0 0 0 0 0 0 1 0
 1 0 1 0 1 0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 1 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0
 0 0 0 0 1 0 0 1 1 0 1 1 0 0 0 0 0 0 0 0 0 0 1 0 1 0 0 0 1 1 0 0 1 1 0 0 1
 1 0 0 0 0 0 0 0 1 0 0 1 1 0 0 0 0 1 1]
[[244.  93.]
 [ 92.  70.]]
[[0.7240356  0.27596438]
 [0.56790125 0.43209878]]
</pre></div>
</div>
</div>
</div>
</div>
</div>
</div>

    <script type="text/x-thebe-config">
    {
        requestKernel: true,
        binderOptions: {
            repo: "binder-examples/jupyter-stacks-datascience",
            ref: "master",
        },
        codeMirrorConfig: {
            theme: "abcdef",
            mode: "python"
        },
        kernelOptions: {
            kernelName: "python3",
            path: "./chapters/hmm"
        },
        predefinedOutput: true
    }
    </script>
    <script>kernelName = 'python3'</script>

              </div>
              
            
                <!-- Previous / next buttons -->
<div class='prev-next-area'>
</div>
            
        </div>
    </div>
    <footer class="footer">
  <p>
    
      By Kevin Murphy, Scott Linderman, et al.<br/>
    
        &copy; Copyright 2021.<br/>
  </p>
</footer>
</main>


      </div>
    </div>
  
  <script src="../../_static/js/index.be7d3bbb2ef33a8344ce.js"></script>

  </body>
</html>