<jupyter_start><jupyter_text>Overview I originally set out to augment the SMoL tool with MAP inference over the branching distribution with a Dirichlet prior, which I did successfully. Additionally, I added MAP inference over the emission distributions of symbols. This second task turned out to be more challenging, since it required rewriting some of the existing SMoL inference engine to work with Automatic Differentiation.This graphic shows an overview of the inference features that existed already in SMoL before this project (left), in contrast to the branching MAP inference (center) and emission MAP inference (right).![alt text](graphical.png "Rough graphical model of SMoL process")Previously when doing inference, I was restricted to viewing sequences as having been generated from the same set of 'decisions', or branching indices. Adding support for inference over the categorical branching distributions allows me to view sequences as having been generated completely separately. To use the genotyping application as an example, I can now model sequences from multiple sources, rather than mutiple reads from the same source and assuming that all variants are the same.Additionally, I can infer emission distributions. This will be useful for another ongoing project.  Inference: MAP Estimates of Dirchlet-Categorical Branch Distributions My strategy for finding the MAP estimate of the branch distribution was to use gradient descent with Haskell's support for automatic differentiation. This not only allows for the use soft-assignment of branch parameters that SMoL produces, but would allow for non-Dirichlet priors to be dropped in. From my existing implementation, I can extract the likelihoods of the data given each possible branch parameter in a model in only one run of the expensive backward-forward algorithm. In order to avoid running the backward-forward algorithm repeatedly, I used the branch likelihoods to intermediate the branch distribution and the complex sequence model. I expressed the posterior probability of the categorical parameters in terms of the branch likelihoods (a dot product, since the branch indices are samples from the categorical), and took the product with the Dirichlet prior to find the joint. I used automatic differentiation to optimize the joint.Below, I demonsrate the use of the new inference capability.<jupyter_code>import SMoL
import SMoL.Inference.SHMM
import SMoL.Tags.Utils
import Control.Monad
import SMoL.DirichletInference
import SMoL.EmissionsInference hiding (simulateEmissionsUniform)
import qualified Data.Map as Map
import Data.Map (Map)
import Data.List (intercalate)
-- Boring helper functions to help with gradient descent convergence and displaying
epsilon = 0.000001
listConverged xs1 xs2 = (< epsilon) . maximum $ zipWith (\x1 x2 -> abs (x1 - x2)) xs1 xs2
mapConverged xs1 xs2 = (< epsilon) . maximum $ Map.intersectionWith (\x1 x2 -> abs (x1 - x2)) xs1 xs2

showMap :: (Show a, Show b) => Map a b -> String
showMap = intercalate "\t" . map (\(k, v) -> show k ++ ":" ++ show v) . Map.toList

showIters :: (a -> String) -> (a -> a -> Bool) -> IO [a] -> IO ()
showIters showF converged = showIters' . (map showF <$>) . (takeUntilPair converged <$>)
  where showIters' :: IO [String] -> IO ()
        showIters' its = its >>= \iters -> forM_ (zip iters [0..]) $ \(iter, n) ->
          putStrLn $ "Iteration " ++ show n ++ ":\t" ++ iter<jupyter_output><empty_output><jupyter_text>Example 1 This is a dead-simple example of branching. Simulated sequences are generated using a categorical distribution over "a", "b" or "c". Data is generated from the sequences by adding noise. The simulated data is modelled with a (not very accurate) Dirichlet prior (implicitely over a categorical distribution). The "runQueries" line runs the optimization using the data, and outputs a sequence of increasingly accurate MAP estimates of the categorical branch distribution. The estimates are printed until they converge.<jupyter_code>-- Sequence model to simulate symbol sequences: p(a)=0.3, p(b)=0.2, p(c)=0.5
generatingSeq = finiteDistOver $ zip [symbol 'a', symbol 'b', symbol 'c'] [0.3, 0.2, 0.5]

-- Prior sequence model, where [2, 5, 3] are dirichlet parameters for the prior of the branch categorical
modelSeq = runTagGen $ finiteDistOverDirichletM [symbol 'a', symbol 'b', symbol 'c'] [2, 5, 3]

-- Sample 1000 example symbol sequences, with noisy emissions distributions
observations = replicateM 1000 (simulateEmissionsUniform "abc" 0.95 <$> sample (compileSMoL generatingSeq))

-- Run the "branchMAPs" query, which runs an MAP optimization over the categorical branch distribution given data
dirichletExample1Iters = runQueries modelSeq branchMAPs <$> observations

-- Print the iteration until convergence
dirichletExample1 = showIters show listConverged dirichletExample1Iters<jupyter_output><empty_output><jupyter_text>We see that the estimates get reasonably close to the truth by the end.  Example 2 This example follows from the branch index example from the other notebook. The goal is to estimate the number of short tandem repeats in DNA, given a signal from a MinION Nanopore Sequencer, which outputs a noisy sliding window. Whereas before we could only estimate the costant number of repeats from the same individual, we can now estimate the distribution of repeats from multiple individuals (the branch index is sampled repeatedly).The true distribution here is generated by "strTruth", and the data is modelled by "strModel [2,2,2,2]", which is a very weak prior Dirichlet.<jupyter_code>-- THIS EXTENDS THE EXAMPLE FROM THE OTHER NOTEBOOK, I RAN THIS BEFORE I SEPARATED THE NOTEBOOKS

-- Repeat the STR segment with the given probabilities (starting with p(0)=0)
strTruth = strProblem (finiteDistRepeat [0.0, 0.15, 0.3, 0.3, 0.25] strSegment)
strTruthC = compileSMoL strTruth

-- Generate an STR model with a given Dirichlet prior, and return the random variable
strModel :: [Double] -> (ProbSeq [NT], DirichletTag)
strModel prior = runTagGen $ do
    (ps, repeatVar) <- finiteDistRepeatDirichletM prior strSegment
    return (strProblem ps, repeatVar)

-- Randomly sample from the true model, with noisy emissions
simulateSTR :: IO (Emissions Double [NT])
simulateSTR = do
    s <- sample strTruthC
    return (simulateEmissionsUniform (minionSymbols k) 0.9 s)

-- Learn the distribution of repeats with a weak prior
dirichletExample2a = showIters show  listConverged $
  runQueries (strModel [2, 2, 2, 2]) branchMAPs <$> replicateM 100 simulateSTR

dirichletExample2a<jupyter_output><empty_output><jupyter_text>We see that the model does.. okay. The next example shows the same setup with a stronger prior.<jupyter_code>-- Learn the distribution of repeats with a stronger prior
dirichletExample2b = showIters show listConverged $
  runQueries (strModel [16, 31, 31, 26]) branchMAPs <$> replicateM 100 simulateSTR
  
dirichletExample2b<jupyter_output><empty_output><jupyter_text>We see that inference is easier with a prior, as expected, and that inference is possible in this odd setting.  Inference: MAP Estimates of Emission Distributions I implemented MAP estimates for emission distribution by rewriting the HMM-like inference algorithm for SMoL. The existing implemenation used an external C++ library for applying the Backward-Forward algorithm while taking advantage of the sparsity of the transition matrix. I replaced this implementation (for the scope of this application) with a pure-Haskell implemenation of the Forward algorithm, from which I can derive the likelihood of the data as a function of the emission parameters. Once I could express the likelihood of the data, I could compute the joint with the prior, and perform the optimization with gradient descent.Rewriting this section in pure, parametric Haskell allowed me to use Haskell's support for automatic differentiation.This section demonstrates learning the emission distributions of symbols, given the usual noisy time-series data and structured sequence models.<jupyter_code>-- Helper function for showing sequence samples
sampleEmissions generatingNorms generatingModel = do
  seq <- fst <$> sampleSeq vecDist (compileSMoL generatingModel)
  putStrLn $ "Generating sequence: " ++ show seq

  ems <- evalState (sampleNormalSequence generatingNorms seq) <$> getStdGen
  putStrLn $ "Generated emissions: " ++ show ems<jupyter_output><empty_output><jupyter_text>Example 1 Up until now, we've been dealing with data where each point is simply a distribution over each possible symbol. Now, we'll view data as a real number per time, and symbols as having a single-variate Gaussian emission.Here's a simple example, where our sequence is "abab", the symbol "a" has emissions from Guassian(-1,1), and the symbol "b" has emissions from Guassian(1,1).<jupyter_code>sampleEmissions (Map.fromList [('a',(-1,1)),('b',(1,1))]) (symbols "abab")<jupyter_output><empty_output><jupyter_text>The "meanDescent" function below is doing the setup for a simulation study and performing the optimization. It takes 100 sequence samples (which are always "abab" for now), and then samples a sequence of real values from those sequences 10 times each (using the distributions from the last argument, with means -1 and 1). It attempts to find the MAP estimate of the emission distribution means, using 0s as starting points and 0-centered means as priors.This is a simple example of emission inference, the MAP estimate does converge to -1 and 1.<jupyter_code>emissionsExample1 = showIters showMap mapConverged $
  meanDescent 100 10 -- sample 100 symbolic sequences, and 10 emissions from each
  (symbols "abab") -- the generating and prior sequence model
  (Map.fromList [('a',0), ('b',0)]) -- the initial distribution means
  (Map.fromList [('a', (0,4)), ('b', (0,4))]) -- the distribution prior distributions (mean, stddev)
  (Map.fromList [('a',-1),('b',1)]) -- the generating emission distribution means

emissionsExample1<jupyter_output><empty_output><jupyter_text>Example 2 In this section, we'll do emission distribution inference with a non-deterministic sequence model:<jupyter_code>exampleEmissions (Map.fromList [('a',(-1,1)),('b',(1,1))]) (eitherOr 0.5 (symbols "ab") (symbols "ba"))<jupyter_output><empty_output><jupyter_text>There is a problem with this model, shown below in the same way as before:<jupyter_code>emissionsExample2a = showIters showMap mapConverged $
  meanDescent 100 10 -- sample 100 symbolic sequences, and 10 emissions from each
  (eitherOr 0.5 (symbols "ab") (symbols "ba")) -- the generating and prior sequence model
  (Map.fromList [('a',0), ('b',0)]) -- the initial distribution means
  (Map.fromList [('a', (0.0,4)), ('b', (0.0,4))]) -- the distribution prior distributions (mean, stddev)
  (Map.fromList [('a',-1),('b',1)]) -- the generating emission distribution means

emissionsExample2a<jupyter_output><empty_output><jupyter_text>The optimization failed to converge because the symbols are symmetrical. This is akin to attempting to fit a mixture of Gaussians clustering with the same starting points.To solve this toy problem, there are two solutions. The first is to introduce asymmetry in the prior distributions:<jupyter_code>emissionsExample2b = showIters showMap mapConverged $
  meanDescent 100 10 -- sample 100 symbolic sequences, and 10 emissions from each
  (eitherOr 0.5 (symbols "ab") (symbols "ba")) -- the generating and prior sequence model
  (Map.fromList [('a',0), ('b',0)]) -- the initial distribution means
  (Map.fromList [('a', (-0.5,4)), ('b', (0.5,4))]) -- the distribution prior distributions (mean, stddev)
  (Map.fromList [('a',-1),('b',1)]) -- the generating emission distribution means
  
emissionsExample2b<jupyter_output><empty_output><jupyter_text>The second is to introduce asymmetry in the sequence model. Since we know "ab" is more likely to occur then "ba", the model learns the difference between the two states.<jupyter_code>emissionsExample2c = showIters showMap mapConverged $
  meanDescent 100 10 -- sample 100 symbolic sequences, and 10 emissions from each
  (eitherOr 0.75 (symbols "ab") (symbols "ba")) -- the generating and prior sequence model
  (Map.fromList [('a',0), ('b',0)]) -- the initial distribution means
  (Map.fromList [('a', (0.0,4)), ('b', (0.0,4))]) -- the distribution prior distributions (mean, stddev)
  (Map.fromList [('a',-1),('b',1)]) -- the generating emission distribution means

emissionsExample2c<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08289283998077847, 'fraction_numerical': 0.01761973410219446, 'mean_word_length': 5.145177165354331, 'pattern_counts': {'":': 2, '<': 59, '<?xml version=': 0, '>': 70, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '292121', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3566, 'n_tokens_neox': 3280, 'n_words': 1572, 'path': 'final-project/final-notebook.ipynb', 'repo_name': 'rybern/smol-shmm', 'size': '12486'}