MainProgram.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fdf8d2fb-47e6-4e07-8435-b51b0607f37e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Info: Precompiling BenchmarkTools [6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf]\n",
      "└ @ Base loading.jl:1662\n"
     ]
    }
   ],
   "source": [
    "import LinearAlgebra\n",
    "import Random\n",
    "import BenchmarkTools\n",
    "import StaticArrays\n",
    "import NBInclude\n",
    "import LinearAlgebra\n",
    "import Distributions\n",
    "import StatsBase\n",
    "import Profile\n",
    "import StatProfilerHTML"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5253d2f8-fe78-42ae-b298-b0c7dd3f5dc9",
   "metadata": {},
   "source": [
    "# sequence simulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4d3bcf62-042f-4624-a79a-f494366b4890",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GenerateSequence (generic function with 1 method)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# this function takes the length as an argument to generate sequences\n",
    "# note, this can also be used to generate a motif\n",
    "function GenerateSequence(Length)\n",
    "    return collect(Random.randstring(\"ACGT\", Length))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "064f10be-856c-4e37-84fd-809c90ac3c8f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GenerateSequences (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create a vector of vector{chars}, each with the same length\n",
    "function GenerateSequences(NumberOfSequences, Length)\n",
    "    return Sequences = map(i -> GenerateSequence(Length), 1:NumberOfSequences)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c0adf295-70ed-405a-9501-57e977f291da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Mutate! (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function Mutate!(Sequence, Distance)\n",
    "    # return the sequence if the Distance is 0\n",
    "    if Distance == 0\n",
    "        return Sequence\n",
    "    end\n",
    "    # choose the sites to mutate\n",
    "    posToMutate = StatsBase.sample(1:length(Sequence), Distance, replace = false)\n",
    "    \n",
    "    # pick the letters that can be used at each position\n",
    "    basesAtPositions = [string(i) for i in Sequence[posToMutate]]\n",
    "\n",
    "    lettersToFill = [\"ACGT\" for i in 1:length(Sequence)]\n",
    "    \n",
    "    # figure out what letters each positions can be changed to\n",
    "    lettersToFill = map((x, y) -> replace(x, y => \"\"), lettersToFill, basesAtPositions)\n",
    "    \n",
    "    # select one character from each of these positions and place them where they should be\n",
    "    # use only to cast the string to a character\n",
    "    # https://stackoverflow.com/questions/59946081/julia-convert-string-to-char-or-convert-arraysubstringstring-1-to-char\n",
    "    Sequence[posToMutate] = map(x -> only(Random.randstring(x, 1)), lettersToFill)\n",
    "    \n",
    "    return Sequence\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f42e095a-babe-4f63-b352-fe32488631cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GenerateTestData_ld (generic function with 1 method)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# generate test data\n",
    "# distance is the number of mutations in each implanted motif\n",
    "function GenerateTestData_ld(NumberOfSequences, LengthMotif, LengthSequences, Distance)\n",
    "    # do some error checking to make sure the values provided are valid\n",
    "    # specifically, Disance <= LengthMotif\n",
    "    if Distance > LengthMotif\n",
    "        error(\"The Distance if larger than the LengthMotif.\")\n",
    "    end\n",
    "\n",
    "    # LengthMotif <= LengthSequences\n",
    "    if LengthMotif > LengthSequences\n",
    "        error(\"The LengthMotif is longer than the LengthSequences\")\n",
    "    end\n",
    "    \n",
    "    # NumberOfSequences >= 2\n",
    "    if NumberOfSequences <= 1\n",
    "        error(\"The NumberOfSequences is 1, which is too small for motif detection\")\n",
    "    end\n",
    "\n",
    "    # generate the input sequences\n",
    "    sequences = GenerateSequences(NumberOfSequences, LengthSequences)\n",
    "    \n",
    "    # make the motif\n",
    "    motif = GenerateSequence(LengthMotif)\n",
    "    \n",
    "    # generate a mutated motif for each sequence to implant\n",
    "    motifs = [copy(motif) for i in 1:NumberOfSequences]\n",
    "    \n",
    "    # broadcast the mutated motif\n",
    "    motifs .= Mutate!.(motifs, Distance)\n",
    "    \n",
    "    # now implant the mutated motifs and record where we implant it\n",
    "    motifStarts = rand(1:(LengthSequences - LengthMotif + 1), NumberOfSequences)\n",
    "    \n",
    "    # now place the motifs into the sequences vector\n",
    "    for i in 1:length(sequences)\n",
    "        sequences[i][motifStarts[i]:(motifStarts[i] + LengthMotif - 1)] = motifs[i]\n",
    "    end\n",
    "    \n",
    "    # now return all the things we may want later\n",
    "    return (motif = motif, motifs_starts = motifStarts, motifs_implanted = motifs, sequences = sequences)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f2aad9-cf7f-4700-ad05-fa0365589945",
   "metadata": {},
   "source": [
    "# Score Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d4e8719b-9ec1-4479-ae23-0a3e26ded086",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Score (generic function with 3 methods)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# it is written in such a way that it is continuously defined\n",
    "function Score(Starts_inp, Length=length(motif), Seqs=sequences; ln=false)\n",
    "    maxScore = Length * length(Seqs)\n",
    "    # MODIFY the inputs if they are out of bounds :-)\n",
    "    ### TO DO: return a very large number of the starts are out of bounds\n",
    "    if @inbounds (minimum(Starts_inp) < 1) | (maximum(Starts_inp) > (length(Seqs[1]) - Length + 1))\n",
    "        return(maxScore)\n",
    "    end\n",
    "    # in the mean time, just fix out of bounds errors\n",
    "    \n",
    "    # generate the character matrix\n",
    "    # then store the results inside of it\n",
    "    #@inbounds @views seqsMatrix = mapreduce(permutedims, vcat, map((s, i) -> s[i:(i + Length - 1)], Seqs, Starts_inp))\n",
    "    @inbounds seqsMatrix = mapreduce(permutedims, vcat, map((s, i) -> s[i:(i + Length - 1)], Seqs, Starts_inp))\n",
    "    \n",
    "    # allocate an array for the most common values in each column\n",
    "    @inbounds mostCommon = Array{Char}(undef, size(seqsMatrix, 2))\n",
    "    thisScore = 0\n",
    "   \n",
    "    # find the most common element in each column \n",
    "    for i in 1:size(seqsMatrix, 2)\n",
    "        A = count(x -> x == 'A', @views @inbounds seqsMatrix[:,i])\n",
    "        C = count(x -> x == 'C', @views @inbounds seqsMatrix[:,i])\n",
    "        G = count(x -> x == 'G', @views @inbounds seqsMatrix[:,i])\n",
    "        T = count(x -> x == 'T', @views @inbounds seqsMatrix[:,i])\n",
    "        \n",
    "        char_max = max(A, C, G, T)\n",
    "        thisScore += char_max\n",
    "        \n",
    "        if char_max == A\n",
    "            @inbounds mostCommon[i] = 'A'\n",
    "        elseif char_max == C\n",
    "            @inbounds mostCommon[i] = 'C'\n",
    "        elseif char_max == G\n",
    "            @inbounds mostCommon[i] = 'G'\n",
    "        else\n",
    "            @inbounds mostCommon[i] = 'T'\n",
    "        end     \n",
    "    end\n",
    "\n",
    "    if ln\n",
    "        return(-1 * log(maxScore - thisScore + 1))\n",
    "    else\n",
    "    # make the minimum (best) score 1 so we can compute the log\n",
    "        return(maxScore - thisScore + 1)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ab587aee-8301-43fc-8d4d-ad6577cfc8f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConsensusMotif (generic function with 3 methods)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# it is written in such a way that it is continuously defined\n",
    "function ConsensusMotif(Starts_inp, Length=length(motif), Seqs=sequences; ln=false)\n",
    "    \n",
    "    # generate the character matrix\n",
    "    # then store the results inside of it\n",
    "    #@inbounds @views seqsMatrix = mapreduce(permutedims, vcat, map((s, i) -> s[i:(i + Length - 1)], Seqs, Starts_inp))\n",
    "    @inbounds seqsMatrix = mapreduce(permutedims, vcat, map((s, i) -> s[i:(i + Length - 1)], Seqs, Starts_inp))\n",
    "    \n",
    "    # allocate an array for the most common values in each column\n",
    "    @inbounds mostCommon = Array{Char}(undef, size(seqsMatrix, 2))\n",
    "    thisScore = 0\n",
    "   \n",
    "    # find the most common element in each column \n",
    "    for i in 1:size(seqsMatrix, 2)\n",
    "        A = count(x -> x == 'A', @views @inbounds seqsMatrix[:,i])\n",
    "        C = count(x -> x == 'C', @views @inbounds seqsMatrix[:,i])\n",
    "        G = count(x -> x == 'G', @views @inbounds seqsMatrix[:,i])\n",
    "        T = count(x -> x == 'T', @views @inbounds seqsMatrix[:,i])\n",
    "        \n",
    "        char_max = max(A, C, G, T)\n",
    "        thisScore += char_max\n",
    "        \n",
    "        if char_max == A\n",
    "            @inbounds mostCommon[i] = 'A'\n",
    "        elseif char_max == C\n",
    "            @inbounds mostCommon[i] = 'C'\n",
    "        elseif char_max == G\n",
    "            @inbounds mostCommon[i] = 'G'\n",
    "        else\n",
    "            @inbounds mostCommon[i] = 'T'\n",
    "        end     \n",
    "    end\n",
    "    \n",
    "    return(mostCommon)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "537610ff-6ba8-443b-b7a9-91678c2a0b66",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Grad (generic function with 1 method)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# calculate a partial score in the plus and minus direction\n",
    "# return positive if going in the positive direction increases the score\n",
    "# return negative if going in the negative direction increases the score\n",
    "# the value given is the sum score difference gained from going x-1 -> x -> x+1\n",
    "#### TO DO: deal with edges of the search space\n",
    "#### TO DO: speed up the update by only updating the score by a change in one dimension\n",
    "# https://www2.atmos.umd.edu/~ekalnay/syllabi/AOSC614/NWP-CH03-2-2.pdf\n",
    "function Grad(Starts_inp, seq, dt, k, Seqs)\n",
    "    backward = copy(Starts_inp)\n",
    "    forward = copy(Starts_inp)\n",
    "    \n",
    "    backward[seq] -= dt\n",
    "    forward[seq] += dt\n",
    "    \n",
    "    #score1 = Score(Starts_inp) - Score(backward)\n",
    "    #score2 = Score(forward) - Score(Starts_inp)\n",
    "    return((Score(forward, k, Seqs; ln = true) - Score(backward, k, Seqs; ln = true))/(2*dt))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d0f53aba-627b-4418-93aa-dca1995fc983",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CalcGrad (generic function with 1 method)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function CalcGrad(StartsInp, dt, k, Seqs)\n",
    "    # make a vector copy we can modify\n",
    "    backward = copy([i for i in StartsInp])\n",
    "    forward = copy([i for i in StartsInp])\n",
    "    grad = zeros(Float64, length(StartsInp))\n",
    "    \n",
    "    for seq in 1:length(StartsInp)\n",
    "        backward[seq] -= dt\n",
    "        forward[seq] += dt\n",
    "        grad[seq] = (Score(forward, k, Seqs; ln = true) - Score(backward, k, Seqs; ln = true))/(2*dt)\n",
    "    end\n",
    "    \n",
    "    return(grad)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d59e4299-4914-4447-a9b3-f6c9ee59225d",
   "metadata": {},
   "source": [
    "# HMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a16b634e-099c-4fcc-a0b1-e9333e45c12f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "H (generic function with 1 method)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function H(particle, velocity, k, Seqs)\n",
    "    return(U(particle, k, Seqs) + 1/2 * velocity' * inv(LinearAlgebra.I) * velocity)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "23c18f28-99af-4a95-b4cf-68dfc232e99a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "U (generic function with 1 method)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function U(particle, k, Seqs)\n",
    "    return(-1 * Score(particle, k, Seqs; ln = true) )\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b83a01c-16a6-4aea-8c18-002fb89dc376",
   "metadata": {},
   "source": [
    "# try to do HMC in place in memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "890c4e20-2216-4011-94f5-53186590f53c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "leap! (generic function with 1 method)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function leap!(particle, velocity, dt, iter, k, Seqs, i, parts, scores)\n",
    "    #### To do: convert these vectors to staticarrays\n",
    "    # initial particle position and velocity\n",
    "    xnt = copy([i for i in particle])\n",
    "    pnt = Float64.(copy(velocity))\n",
    "    pn_t_dt = zeros(Float64, length(particle))\n",
    "    pn_t_dt2 = zeros(Float64, length(particle))\n",
    "    xn_t_dt = copy([i for i in particle])\n",
    "    \n",
    "    # iterate for the number of time steps desired\n",
    "    # TODO: add stopping criterion\n",
    "    for t in 1:iter\n",
    "        # do the first half-step\n",
    "        # update the particles momentum by looking at the previous moments\n",
    "        # and adding half-ish the change in the potential energy\n",
    "        pn_t_dt2 .= pnt .- (dt/2) .* CalcGrad(xnt, dt, k, Seqs)\n",
    "\n",
    "        # update the particle position after t time\n",
    "        # by using the new momentum, found in the previous half step\n",
    "        \n",
    "        # round this value to an integer, since our scorer only works\n",
    "        # on discrete values\n",
    "        ### TODO: add a bounds check here to prevent from going out of bounds\n",
    "        # also, do something different besides rounding, like using the logic in add_vectors\n",
    "        xn_t_dt .= Int.(round.(xnt .+ dt * inv(LinearAlgebra.I) * pn_t_dt2))\n",
    "        \n",
    "        # update the momentum again, this time by evaluating the gradient at the new point\n",
    "        pn_t_dt .= pn_t_dt2 .- (dt/2) .* CalcGrad(xn_t_dt, dt, k, Seqs)\n",
    "        \n",
    "        # do a metropolis-hastings step\n",
    "        \n",
    "        # compute the ratio of hamiltonian new/old like on the wiki page      \n",
    "        # if this value is off substantially, we will end up rejecting the sample\n",
    "        α = min(1, (exp(-1 * H(xn_t_dt, pn_t_dt, k, Seqs)))/(exp(-1 * H(xnt, pnt, k, Seqs))))\n",
    "        # quit if we're deviating substantially from the HMC path\n",
    "        if α < 0.0001\n",
    "            return \n",
    "        end\n",
    "        if (minimum(xn_t_dt) < 1) | (maximum(xn_t_dt) > (length(Seqs[1]) - k + 1))\n",
    "            break\n",
    "        end\n",
    "        \n",
    "        # if our random number is less than alpha, accept the new proposal\n",
    "        if rand(1)[1] <= α\n",
    "            xnt .= xn_t_dt\n",
    "            pnt .= pn_t_dt\n",
    "            parts[i[1]] = StaticArrays.SVector{length(particle), Int64}(xnt)\n",
    "            scores[i[1]] = Score(xnt, k, Seqs)\n",
    "            i .+= 1\n",
    "        end\n",
    "        \n",
    "        #### TODO:\n",
    "        # check if we are at a boundary. if so, reverse the sign of the velocity for that dimension\n",
    "    end\n",
    "    #### TODO: also return the scores of each particle\n",
    "    return\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c3400c8f-b3ed-4498-a95d-b75ee072ddaf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "make_sample! (generic function with 1 method)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function make_sample!(particle, velocity, num, dt, iter, k, Seqs, parts, scores)\n",
    "    # TODO: preallocate the vector for the particles\n",
    "    # likely need these two vectors on from Algorithm, and just keep reusing the same vectors over and over\n",
    "    # preallocate scores with a very large number probably, maybe the largest Int64?\n",
    "    this_part = zeros(Int64, length(particle))\n",
    "    this_velo = zeros(Int64, length(velocity))\n",
    "    i = [1]\n",
    "    # pre allocate the vector for the sample\n",
    "    for n in 1:num\n",
    "        this_part .= particle .+ velocity .+ StatsBase.sample(-5:5, length(particle)) ### TODO: initiate the starting position of the particle based on particle and velocity\n",
    "        this_velo .= StatsBase.sample(-3:3, length(particle)) ### TODO: initiate the velocity\n",
    "        # sample random starting points at particle + \n",
    "        leap!(this_part, this_velo, dt, iter, k, Seqs, i, parts, scores)\n",
    "    end\n",
    "    \n",
    "    # handle the case where no sampling is done\n",
    "    if length(scores) == 0\n",
    "        return(parts, [], Inf)\n",
    "    end\n",
    "    minpos = argmin(scores)\n",
    "    bestpart = parts[minpos]\n",
    "    bestscore = scores[minpos]\n",
    "    # return the sample, as well as the (particle, score) having the best score\n",
    "    return(bestpart, bestscore)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8d39c28d-217a-4ba2-9172-c16c6b40b836",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "900-element Vector{Int64}:\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       "                   ⋮\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807\n",
       " 9223372036854775807"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parts = [StaticArrays.SVector{3, Int64}(zeros(3)) for i in 1:(30*30)]\n",
    "scores = [typemax(Int64) for i in 1:(30*30)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5385f0a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "900-element Vector{StaticArraysCore.SVector{3, Int64}}:\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " ⋮\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]\n",
       " [0, 0, 0]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9583265e-0733-49d6-b4ce-dd2c34b995de",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3068708"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@allocated parts = [StaticArrays.SVector{3, Int64}(zeros(3)) for i in 1:(30*30)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b585cc49-9b61-4497-8535-b7e16f88c5b5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "175504"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@allocated begin\n",
    "    for i in 1:length(parts)\n",
    "        parts[i] = StaticArrays.SVector{3, Int64}(zeros(3))\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bce70b6-a284-4aaf-b102-1034c01e9303",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Random.seed!(123)\n",
    "#@allocated make_sample!(StaticArrays.SVector{3, Int64}(10, 10, 10), StaticArrays.SVector{3, Int64}(1, 1, 1), 30, 1, 30, 10, sequences, parts, scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f93c0bdc-8433-43d2-9e8d-896ba0adbc5f",
   "metadata": {
    "tags": []
   },
   "source": [
    "# PSO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fc10904-fe46-4171-8023-1622117466db",
   "metadata": {},
   "source": [
    "# try some different logic in add_vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "527fd7ee-c94a-483d-a934-1e5c7ebed4e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "add_vectors (generic function with 1 method)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function add_vectors(vec1, vec2, energy)\n",
    "    # if vector to add is all zeros, just return the original vector\n",
    "    if vec2 == zeros(length(vec1))\n",
    "        return vec1\n",
    "    end\n",
    "\n",
    "    # find the norm of the first vector, which is the particle\n",
    "    #added = (vec1 .+ vec2).*(LinearAlgebra.norm(vec1, 2)/(LinearAlgebra.norm(vec1, 2)+LinearAlgebra.norm(vec2, 2)))\n",
    "    \n",
    "    # current_sign = sign.(added)\n",
    "    \n",
    "    # get the remainder \n",
    "    # rounded_ver = abs.(added) .- floor.(abs.(added))\n",
    "    \n",
    "    # for the negative numbers, take 1 minus the value\n",
    "    # rounded_ver[current_sign .== -1.0] .= 1 .- rounded_ver[current_sign .== -1.0]\n",
    "\n",
    "    # get the remainder part\n",
    "    # round_dir = ifelse.(rounded_ver .< rand(length(vec1)), RoundUp::RoundingMode, RoundDown::RoundingMode)\n",
    "    # add the two vectors together\n",
    "    # now round each value to get an integer array\n",
    "    # get a random number for each position in the vector\n",
    "    \n",
    "    # round the vector\n",
    "    # vec_temp = map((x, y) -> round(Int64, x, y), added, round_dir)\n",
    "    vec_temp = vec1 .+ vec2\n",
    "    norm_temp = LinearAlgebra.norm(vec_temp)\n",
    "    \n",
    "    # make the vector longer if the norm is smaller than the energy\n",
    "    if norm_temp < energy\n",
    "        while (LinearAlgebra.norm(vec_temp) < energy) & (LinearAlgebra.norm(vec_temp) > 0)\n",
    "            # pick a random position in the vector\n",
    "            rand_pos = rand(1:length(vec_temp))\n",
    "            # now subtract its sign from that position\n",
    "            vec_temp[rand_pos] += sign(vec_temp[rand_pos])\n",
    "        end\n",
    "    # make the vector shorter if the norm is larger than the energy\n",
    "    elseif norm_temp > energy\n",
    "        while (LinearAlgebra.norm(vec_temp) > energy) & (LinearAlgebra.norm(vec_temp) > 0)\n",
    "            # pick a random position in the vector\n",
    "            rand_pos = rand(1:length(vec_temp))\n",
    "            # now add its sign from that position\n",
    "            vec_temp[rand_pos] -= sign(vec_temp[rand_pos])\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    # if any of the dimensions are zeros - try adding 1\n",
    "    \n",
    "    vec_out = StaticArrays.SVector{length(vec1), Int64}(vec_temp)\n",
    "    \n",
    "    # if all equal to zero, just use the sign of vec2\n",
    "    if vec_out == zeros(length(vec1))\n",
    "        return sign.(vec2) .* StatsBase.sample(1:3, length(vec1))\n",
    "    else\n",
    "        return vec_out\n",
    "    end\n",
    "    ## todo: deal with situations where the norm of vec1 is very close tozero\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c4404f2f-1ede-45ec-80f2-ea9613b39c72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element StaticArraysCore.SVector{3, Int64} with indices SOneTo(3):\n",
       " 0\n",
       " 4\n",
       " 5"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "add_vectors([4, 4, 4], [-2, 1, 2], 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3c83698c-6825-4e83-8ae6-6c7b8107ddba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10-element Vector{Int64}:\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "GetFinalParticle(sequences, 10, ['G', 'A', 'T', 'C', 'C', 'T', 'G', 'G', 'T', 'A'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "12e7e486",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10-element Vector{Vector{Char}}:\n",
       " ['C', 'C', 'T', 'T', 'C', 'A', 'G', 'G', 'C', 'T']\n",
       " ['A', 'G', 'C', 'C', 'G', 'T', 'G', 'G', 'G', 'A']\n",
       " ['C', 'C', 'T', 'A', 'C', 'T', 'T', 'A', 'T', 'G']\n",
       " ['A', 'C', 'A', 'G', 'C', 'C', 'A', 'A', 'T', 'G']\n",
       " ['C', 'T', 'A', 'A', 'A', 'C', 'G', 'A', 'G', 'C']\n",
       " ['T', 'T', 'T', 'T', 'A', 'A', 'C', 'T', 'A', 'T']\n",
       " ['C', 'G', 'G', 'T', 'C', 'C', 'G', 'G', 'G', 'A']\n",
       " ['T', 'G', 'C', 'G', 'G', 'G', 'C', 'C', 'A', 'T']\n",
       " ['C', 'T', 'G', 'G', 'G', 'C', 'A', 'C', 'G', 'C']\n",
       " ['G', 'C', 'A', 'G', 'T', 'T', 'T', 'A', 'A', 'T']"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sequences = GenerateSequences(10,10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0f124d58-52b4-4c3c-9529-229281f1fb7b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GetFinalParticle (generic function with 1 method)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function GetFinalParticle(seqs, k, consensus)\n",
    "    # go over all the sequences\n",
    "    out_vec = zeros(Int64, length(seqs))\n",
    "    \n",
    "    for i in 1:length(sequences)\n",
    "        # score the starting positions in each sequence\n",
    "        dist = map(start -> sum(consensus .== seqs[i][start:(start + k - 1)]), 1:(length(seqs[i]) - k + 1))\n",
    "        out_vec[i] = argmax(dist)\n",
    "    end\n",
    "    return(out_vec)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bcec8cfe-ed7d-4067-808c-8391ed76e9cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Algorithm (generic function with 4 methods)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function Algorithm(seqs, k, dims = 3, nparticles = 10, maxiter = 500)\n",
    "    debug = false\n",
    "    # generate the starting points for the algorithm\n",
    "    # store these in a matrix, where each column is a particle, and each row is a dimension in the search space\n",
    "    particles = [StaticArrays.SVector{dims, Int64}(rand(1:(length(seqs[1]) - k + 1), dims)) for i in 1:nparticles]\n",
    "    velocities = [StaticArrays.SVector{dims, Int64}(rand(-5:5, dims)) for i in 1:nparticles]\n",
    "    energies = Int.(round.([LinearAlgebra.norm(i) for i in velocities]))\n",
    "\n",
    "    scores = zeros(Int64, nparticles)\n",
    "    \n",
    "    # create a set to store the points that we have visited, to avoid looking at the same points repeatedly\n",
    "    visited = Set{StaticArrays.SVector{dims, Int64}}()\n",
    "    \n",
    "    # give a size hint for the set of visited points, hopefully preventing it from growing repeatedly\n",
    "    sizehint!(visited, nparticles * maxiter)\n",
    "    \n",
    "    iter = 1\n",
    "    \n",
    "    best_particle = zeros(Int64, dims)\n",
    "    global_optimum = Inf\n",
    "\n",
    "    terminated = 0\n",
    "    \n",
    "    # reuse the same value for sign over and over\n",
    "    sign_vec = zeros(Int64, dims)\n",
    "    \n",
    "    # create the vectors for storing the samples\n",
    "    parts = [StaticArrays.SVector{dims, Int64}(zeros(dims)) for i in 1:(30*30)]\n",
    "    scores_samp = [typemax(Int64) for i in 1:(30*30)]\n",
    "    idx = 1\n",
    "    # loop until convergence\n",
    "    \n",
    "    if debug\n",
    "        for p in 1:nparticles\n",
    "            println(particles[p], velocities[p])\n",
    "        end\n",
    "\n",
    "        println()\n",
    "    end\n",
    "    \n",
    "    while iter <= maxiter\n",
    "        # wrap particles around that have gone out of bounds\n",
    "        counter = 0\n",
    "        # \"wiggle\" each particle position if it's already been visited, or if its out of bounds\n",
    "        # for each of these, try wiggling\n",
    "        for part in findall(x -> (x in visited) | (minimum(x) < 1) | (maximum(x) > (length(seqs[1]) - k + 1)), particles)\n",
    "            counter += 1\n",
    "            newpart = copy(particles[part])\n",
    "            tries = 1\n",
    "            while (tries < 40) & (newpart in visited)\n",
    "                dist = Int(floor(tries/4 + 1))\n",
    "                newpart = particles[part] .+ (rand(1:dist, dims) .* rand([-1, 0, 1], dims))\n",
    "                tries += 1\n",
    "            end\n",
    "            particles[part] = newpart\n",
    "        end\n",
    "        \n",
    "\n",
    "        # compute the current scores of each particle\n",
    "        #scores .= map(x -> Score(x, k, seqs), particles)\n",
    "        scores .= Score.(particles, (k,), (seqs,))\n",
    "        \n",
    "        # find the current best particle\n",
    "        best_score = findmin(scores)\n",
    " \n",
    "        # if this score is better than the current best score, update\n",
    "        if best_score[1] < global_optimum\n",
    "            global_optimum = copy(best_score[1])\n",
    "            best_particle .= copy(particles[best_score[2]])\n",
    "        end\n",
    "    \n",
    "        # for each of the current vectors\n",
    "            # sample nearby positions and run HMC for 30 steps for each of 30 sampled particles\n",
    "            # sample such that particles are more likely to start in the direction the particle is heading and in the direction of the best particle\n",
    "\n",
    "        for i in 1:length(particles)\n",
    "            # make sure the samples are set to be empty\n",
    "            for i in 1:idx\n",
    "                parts[i] = StaticArrays.SVector{dims, Int64}(zeros(dims))\n",
    "                scores_samp[i] = typemax(Int64)\n",
    "            end\n",
    "            \n",
    "            # try adding the direction rather than the sign\n",
    "            sign_vec .= sign.(best_particle .- particles[i])\n",
    "            # sign_vec .= best_particle .- velocities[i]\n",
    "            # update the velocity first to be more in the direction of the best particle\n",
    "            # keep constant energy/total velocity norm for each particle\n",
    "            velocities[i] = add_vectors(velocities[i], sign_vec, energies[i])\n",
    "\n",
    "            # draw the sample\n",
    "            # make_sample!(particle, velocity, num, dt, iter, k, Seqs, parts, scores)\n",
    "            bestpart_sample, bestscore_sample = make_sample!(particles[i], velocities[i], 30, 1, 30, k, seqs, parts, scores_samp)\n",
    "            idx = findfirst(isequal(zeros(dims)), parts)\n",
    "            \n",
    "            # check the sample to see if we found a really good position to aim towards\n",
    "            if bestscore_sample < global_optimum\n",
    "                global_optimum = copy(bestscore_sample)\n",
    "                best_particle .= copy(bestpart_sample)\n",
    "            end\n",
    "\n",
    "            # handle the case where sampling fails since we're on the edge of the search space\n",
    "            if idx > 1\n",
    "                # pick a random term from the sample and update the particle position\n",
    "                @inbounds particles[i] = StatsBase.sample(@views parts[1:(idx - 1)])\n",
    "                # now update to be in the direction of the best particle\n",
    "                \n",
    "            end\n",
    "        end        \n",
    "        \n",
    "        # save the new locations we've visited\n",
    "        map(x -> push!(visited, x), particles)\n",
    "\n",
    "        # update the particle positions based on the velocities\n",
    "        \n",
    "        # convergence criteria: quit when idx is full for ten iterations in a row\n",
    "        if counter == nparticles\n",
    "            terminated += 1\n",
    "        end\n",
    "        \n",
    "        if terminated == 10\n",
    "            iter += maxiter\n",
    "        end\n",
    "\n",
    "        iter += 1\n",
    "        \n",
    "        if debug\n",
    "            for p in 1:nparticles\n",
    "                println(particles[p], velocities[p], Score(particles[p], k, seqs))\n",
    "            end\n",
    "            println(global_optimum, best_particle)\n",
    "            println()\n",
    "        end\n",
    "    end \n",
    "    \n",
    "    # may want to repeat this step multiple times to refine motif\n",
    "    consensus_motif = ConsensusMotif(best_particle, k, seqs)\n",
    "    final_particle = GetFinalParticle(seqs, k, consensus_motif)\n",
    "    final_score = Score(final_particle, k, seqs)\n",
    "    \n",
    "    consensus_motif = ConsensusMotif(final_particle, k, seqs)\n",
    "    \n",
    "    # add a shift step to see if moving up or down helps the score\n",
    "    ## finishing step\n",
    "    ## calculate the position weight matrix, and find the highest scoring motif in each sequence.\n",
    "    \n",
    "    \n",
    "    # also return the consensus motif\n",
    "    return(final_score, final_particle, consensus_motif)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cba76d53-0e1e-4287-a5ce-c45611ad3b7d",
   "metadata": {},
   "source": [
    "# Helpful Output formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "60a27235-de39-477e-a243-c578dc8cca78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "print_motifs (generic function with 1 method)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function print_motifs(Starts_Inp, Seqs, k)\n",
    "    for i in 1:length(Starts_Inp)\n",
    "        println(join(Seqs[i][(Starts_Inp[i]):(Starts_Inp[i] + k - 1)]))\n",
    "    end\n",
    "    return\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4135f980-1b6a-4b33-84bd-7e7d52ed6c57",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Unit Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9e0e52-5540-4e83-8893-59c208c2483f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test add_vectors!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969b81dc-f412-434b-bf72-48ac0919a284",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if out of bounds is working in Score\n",
    "# motif, motif_starts, motifs_implanted, sequences = GenerateTestData_ld(3, 14, 30, 0)\n",
    "# Score([1, 1, 0], sequences, 14)\n",
    "# Score([1, 1, 18], sequences, 14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c220776f-e957-4fa6-b194-1a0bfa7e855f",
   "metadata": {},
   "source": [
    "# Hyperparameter Tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a4d2ed9-1e76-4b32-a25a-1f4451db8450",
   "metadata": {},
   "source": [
    "# Benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d5265a9-14e3-46f7-a3d0-c853a07c1747",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the seed\n",
    "Random.seed!(101)\n",
    "# use tuple unpacking to get some test values\n",
    "# NumberOfSequences, LengthMotif, LengthSequences, Distance\n",
    "motif, motif_starts, motifs_implanted, sequences = GenerateTestData_ld(3, 10, 30, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66023024-2c26-49a2-9e40-05b816e50ae6",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_vec = Algorithm(sequences, 10, size(sequences, 1), 3, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "472ef73d-5cd7-48eb-99a3-fcb401e21251",
   "metadata": {},
   "outputs": [],
   "source": [
    "Random.seed!(123)\n",
    "@allocated Algorithm(sequences, 10, size(sequences, 1), 10, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b28d6e0-c394-4329-86e9-c4c9334cbbcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "BenchmarkTools.@btime Algorithm(sequences, 10, size(sequences, 1), 30, 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba8ac7e-6eee-4b2d-9e32-7d63b3ac54c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "StatProfilerHTML.@profilehtml Algorithm(sequences, 10, size(sequences, 1), 30, 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e3f870f-6b42-4f4e-b499-5edf041b9e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the seed\n",
    "Random.seed!(100)\n",
    "# use tuple unpacking to get some test values\n",
    "# NumberOfSequences, LengthMotif, LengthSequences, Distance\n",
    "motif, motif_starts, motifs_implanted, sequences = GenerateTestData_ld(50, 10, 250, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330441d6-60ff-45b3-8878-3a565c0f9b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "StatProfilerHTML.@profilehtml Algorithm(sequences, 10, size(sequences, 1), 30, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ad207c2-dccf-4172-8ad2-99fa928a795b",
   "metadata": {},
   "source": [
    "# Visualizing HMC sampling in 2D"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.8.2",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}