text:
# # Solve the Stokes equation of colliding flow: Reddy formulation

# Synopsis: Compute the solution of the Stokes equation of two-dimensional
# incompressible viscous flow for a manufactured problem of colliding flow.
# Hood-Taylor triangular elements are used.

# The "manufactured" colliding flow example from Elman et al 2014. The
# Hood-Taylor formulation with quadratic triangles for the velocity and
# continuous pressure on linear triangles.

# The pressure is shown here with contours, and the velocities visualized with
# arrows at random points.
# ![Pressure and velocity](colliding.png)

# The formulation is the one derived in Reddy, Introduction to the finite
# element method, 1993. Page 486 ff.

# The complete code is in the file [`tut_stokes_ht_p2_p1_reddy.jl`](tut_stokes_ht_p2_p1_reddy.jl).

# The solution will be defined  within a module in order to eliminate conflicts
# with data or functions defined elsewhere.

module tut_stokes_ht_p2_p1_reddy

# We'll need some functionality from linear algebra, static arrays, and the mesh
# libraries. Some plotting will be produced to visualize structure of the
# stiffness matrix. Finally we will need the `Elfel` functionality.
using LinearAlgebra
using StaticArrays
using MeshCore.Exports
using MeshSteward.Exports
using Elfel.Exports
using UnicodePlots

# The boundary value problem is expressed in this weak form
# ```math
# \int_{\Omega} \left(2\mu\frac{\partial{w_x}}{\partial{x}}\frac{\partial{u_x}}{\partial{x}}
# +\mu\frac{\partial{w_x}}{\partial{y}}\left(\frac{\partial{u_x}}{\partial{y}}+\frac{\partial{u_y}}{\partial{x}}\right)  - \frac{\partial{w_x}}{\partial{x}} p \right) d\Omega  = 0
# ```
# ```math
# \int_{\Omega} \left(2\mu\frac{\partial{w_y}}{\partial{y}}\frac{\partial{u_y}}{\partial{y}}
# +\mu\frac{\partial{w_y}}{\partial{x}}\left(\frac{\partial{u_x}}{\partial{y}}+\frac{\partial{u_y}}{\partial{x}}\right)  - \frac{\partial{w_y}}{\partial{y}} p \right) d\Omega  = 0
# ```
# ```math
# -\int_{\Omega} q \left(\frac{\partial{u_x}}{\partial{x}} + \frac{\partial{u_y}}{\partial{y}}\right)d\Omega = 0
# ```

# Here ``w_x, w_y`` are the test functions in the velocity space,
# and ``q`` is the pressure test functions. Further ``u_x, u_y`` is the trial velocity, and ``p`` is the trial pressure.

function run()
    mu = 1.0 # dynamic viscosity
    A = 1.0 # half of the length of the side of the square
    N = 100 # number of element edges per side of the square
    # These three functions define the true velocity components and the true
    # pressure.
    trueux = (x, y) -> 20 * x * y ^ 3
    trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
    truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

    # Construct the two meshes for the mixed method. They need to support the
    # velocity and pressure spaces.
    vmesh, pmesh = genmesh(A, N)

    # Construct the velocity spaces: As an alternative to a previous treatment
    # with a single vector space for the velocity, here we will use to vector
    # spaces, one for each component of the velocity. The degrees of freedom
    # are real numbers (`Float64`). The velocity mesh carries the finite
    # elements of  the continuity ``H ^1``, i. e. both the function values and
    # the derivatives are square integrable. Each node carries just one degree of
    # freedom (1).
    Uxh = FESpace(Float64, vmesh, FEH1_T6(), 1)
    Uyh = FESpace(Float64, vmesh, FEH1_T6(), 1)

    # Now we apply the boundary conditions at the nodes around the
    # circumference.
    locs = geometry(vmesh)
    # We use searching based on the presence of the node within a box. The
    # entire boundary will be captured within these four boxes, provided we
    # inflate those boxes with a little tolerance (we can't rely on those
    # nodes to be precisely at the coordinates given, we need to introduce some
    # tolerance).
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    inflate = A / N / 100
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            # Remember that all  components of the velocity are known at the
            # boundary.
            setebc!(Uxh, 0, i, 1, trueux(locs[i]...))
            setebc!(Uyh, 0, i, 1, trueuy(locs[i]...))
        end
    end

    # No we construct the pressure space. It is a continuous, piecewise linear
    # space supported on a mesh of three-node triangles.
    Ph = FESpace(Float64, pmesh, FEH1_T3(), 1)

    # The pressure in this "enclosed" flow example is only known up to a constant.
    # By setting  pressure degree of freedom at one node will make the solution
    # unique.
    atcenter = vselect(geometry(pmesh); nearestto = [0.0, 0.0])
    setebc!(Ph, 0, atcenter[1], 1, 0.0)

    # Number the degrees of freedom. First all the free degrees of freedom are
    # numbered, both velocities and pressures. Next all the data degrees of
    # freedom are numbered, again both for the velocities and for the
    # pressures.
    numberdofs!([Uxh, Uyh, Ph])
    # The total number of degrees of freedom is now calculated.
    tndof = ndofs(Uxh) + ndofs(Uyh) + ndofs(Ph)
    # As is the total number of unknowns.
    tnunk = nunknowns(Uxh) + nunknowns(Uyh) + nunknowns(Ph)

    # Assemble the coefficient matrix.
    K = assembleK(Uxh, Uyh, Ph, tndof, mu)

    # Display the structure of the indefinite stiffness matrix. Note that this
    # is the complete matrix, including rows and columns for all the degrees of
    # freedom, unknown and known.
    p = spy(K, canvas = DotCanvas)
    display(p)

    # Solve the linear algebraic system. First construct system vector of all
    # the degrees of freedom, in the first `tnunk` rows that corresponds to the
    # unknowns, and the subsequent rows are for the data degrees of freedom.
    U = fill(0.0, tndof)
    gathersysvec!(U, [Uxh, Uyh, Ph])
    # Note that  the vector `U` consists of nonzero numbers in rows are for the
    # data degrees of freedom. Multiplying the stiffness matrix with this
    # vector will generate a load vector  on the right-hand side. Otherwise
    # there is no loading, hence the vector `F` consists of all zeros.
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    # Once we have solved the system of linear equations, we can distribute the
    # solution from the vector `U` into the finite element spaces.
    scattersysvec!([Uxh, Uyh, Ph], U)

    # Given that the solution is manufactured, i. e. exactly known, we can
    # calculate the true errors.
    @show ep = evaluate_pressure_error(Ph, truep)
    @show ev = evaluate_velocity_error(Uxh, Uyh, trueux, trueuy)

    # Postprocessing. First we make attributes, scalar nodal attributes,
    # associated with the meshes for the pressures and the velocity.
    makeattribute(Ph, "p", 1)
    makeattribute(Uxh, "ux", 1)
    makeattribute(Uyh, "uy", 1)
    # The pressure and the velocity components are then written out into two VTK
    # files.
    vtkwrite("tut_stokes_ht_p2_p1_reddy-p", baseincrel(pmesh), [(name = "p",), ])
    vtkwrite("tut_stokes_ht_p2_p1_reddy-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])

    return true
end

function genmesh(A, N)
    # Hood-Taylor pair of meshes is needed. The first mesh is for the
    # velocities, composed of six-node triangles.
    vmesh = attach!(Mesh(), T6block(2 * A, 2 * A, N, N), "velocity")
    # Now translate so that the center of the square is at the origin of the
    # coordinates.
    ir = baseincrel(vmesh)
    transform(ir, x -> x .- A)
    # The second mesh is used for the pressures, and it is composed of
    # three-node triangles such that the corner nodes are shared between the
    # first and the second mesh.
    pmesh = attach!(Mesh(), T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    # Return the pair of meshes
    return vmesh, pmesh
end

function assembleK(Uxh, Uyh, Ph, tndof, mu)
    # Here we demonstrate that the coefficient matrix, which is expected to have
    # the structure
    # ```math
    # K = \left[
    # \begin{array}{cc}
    #     A & B^T \\
    #     B & 0
    # \end{array}\right]
    # ```
    # can be constructed in stages. Refer to the description below. The two
    # functions below  carry out the integration of two separate parts of the
    # coefficient matrix.
    function integrateApart!(ass, elits, qpits, mu)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxux = LocalMatrixAssembler(uxnedof, uxnedof, 0.0)
        kuyuy = LocalMatrixAssembler(uynedof, uynedof, 0.0)
        kuxuy = LocalMatrixAssembler(uxnedof, uynedof, 0.0)
        for el in zip(elits...)
            uxel, uyel, pel = el
            init!(kuxux, eldofs(uxel), eldofs(uxel))
            init!(kuyuy, eldofs(uyel), eldofs(uyel))
            init!(kuxuy, eldofs(uxel), eldofs(uyel))
            for qp in zip(qpits...)
                # Step the quadrature point iterators in step: this assumes that
                # in fact there is the same number of quadrature points in all
                # the quadrature rules.
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(pel, pqp)
                JxW = J * weight(pqp)
                # Note that the gradients of the basis functions are not
                # necessarily the same in those two velocity spaces. Hence we
                # must grab the gradient from the correct space.
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                for j in 1:uxnedof, i in 1:uxnedof
                    kuxux[i, j] += (mu * JxW) * (2 * gradNux[i][1] * gradNux[j][1] + gradNux[i][2] * gradNux[j][2])
                end
                for j in 1:uynedof, i in 1:uynedof
                    kuyuy[i, j] += (mu * JxW) * (gradNuy[i][1] * gradNuy[j][1] + 2 * gradNuy[i][2] * gradNuy[j][2])
                end
                for j in 1:uynedof, i in 1:uxnedof
                    kuxuy[i, j] += (mu * JxW) * (gradNux[i][1] * gradNuy[j][2])
                end
            end
            assemble!(ass, kuxux)
            assemble!(ass, kuxuy) # off-diagonal matrix needs to be assembled twice
            assemble!(ass, transpose(kuxuy))
            assemble!(ass, kuyuy)
        end
        return ass
    end
    function integrateBBTparts!(ass, elits, qpits)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxp = LocalMatrixAssembler(uxnedof, pnedof, 0.0)
        kuyp = LocalMatrixAssembler(uynedof, pnedof, 0.0)
        for el in zip(elits...)
            # The iterators of the finite elements are stepped in unison.
            uxel, uyel, pel = el
            # Initialize the local matrix assembler with the global degree of
            # freedom numbers, both for the velocity spaces and for the
            # pressure space.
            init!(kuxp, eldofs(uxel), eldofs(pel))
            init!(kuyp, eldofs(uyel), eldofs(pel))
            for qp in zip(qpits...)
                # Step the quadrature point iterators in step: this assumes that
                # in fact there is the same number of quadrature points in all
                # the quadrature rules.
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(pel, pqp)
                JxW = J * weight(pqp)
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                Np = bfun(pqp)
                for j in 1:pnedof, i in 1:uxnedof
                    kuxp[i, j] += (-JxW) * (gradNux[i][1] * Np[j])
                end
                for j in 1:pnedof, i in 1:uynedof
                    kuyp[i, j] += (-JxW) * (gradNuy[i][2] * Np[j])
                end
            end
            assemble!(ass, kuxp)
            assemble!(ass, transpose(kuxp))
            assemble!(ass, kuyp)
            assemble!(ass, transpose(kuyp))
        end
        return ass
    end

    # In the `assembleK` function we first we create the element iterators. We
    # can go through all the elements, both in the velocity finite element
    # space and in the pressure finite element space, that define the domain of
    # integration using this iterator. Each time a new element is accessed,
    # some data are precomputed such as the element degrees of freedom,
    # components of the degree of freedom, etc. Note that we need to iterate
    # two finite element spaces, hence we create a tuple of iterators.
    elits = (FEIterator(Uxh), FEIterator(Uyh), FEIterator(Ph))
    # These are the quadrature point iterators. We know that the elements are
    # triangular. We choose the three-point rule, to capture the quadratic
    # component in the velocity space. Quadrature-point iterators provide
    # access to basis function values and gradients, the Jacobian matrix and
    # the Jacobian determinant, the location of the quadrature point and so on.
    # Note that we need to iterate the quadrature rules of three finite element
    # spaces, hence we create a tuple of iterators. All of these quadrature
    # point iterators refer to the "same" quadrature rule: the same number of
    # quadrature points, the same weights, and so on. However, the services
    # these quadrature points provide do depend on the finite element space as
    # well, for instance they would typically have different basis functions.
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(Uxh, qargs), QPIterator(Uyh, qargs), QPIterator(Ph, qargs))
    # The matrix will be assembled into this assembler. Which is initialized
    # with the total number of degrees of freedom (dimension of the coefficient
    # matrix before partitioning into unknowns and data degrees of freedom).
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    # First we calculate the "A" part, using the function below. It
    # is "assembled" into the assembler, which means that after this function
    # finishes, the assembler represents this intermediate matrix
    # ```math
    # K = \left[
    # \begin{array}{cc}
    #     A & 0 \\
    #     0 & 0
    # \end{array}\right]
    # ```
    integrateApart!(ass, elits, qpits, mu)
    # Then the "B-transpose(B)" part using this function is added to the
    # assembler. When the function below finishes, the assembler represents the
    # entire ``K`` matrix.
    integrateBBTparts!(ass, elits, qpits)
    # Finally, we materialize the sparse stiffness matrix from the assembler and
    # return it.
    return finish!(ass)
end

# The linear algebraic system is solved by partitioning. The vector `U` is
# initially all zero, except in the degrees of freedom which are prescribed as
# nonzero. Therefore the product of the stiffness matrix and the vector `U`
# are the loads due to nonzero essential boundary conditions.  The
# submatrix of the stiffness conduction matrix corresponding to the free degrees of
# freedom (unknowns), `K[1:nu, 1:nu]` is then used to solve for the unknowns `U
# [1:nu]`.
function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

# The function `evaluate_pressure_error` evaluates the true ``L^2``  error of
# the pressure. It does that by integrating the square of the difference
# between the approximate pressure and the  true pressure, the true pressure
# being provided by the `truep` function.
function evaluate_pressure_error(Ph, truep)
    function integrate!(elit, qpit, truep)
        n_dp = ndofsperel(elit)
        E = 0.0
        for el in elit
            dofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Np = bfun(qp)
                pt = truep(location(el, qp)...)
                pa = 0.0
                for j in 1:n_dp
                    pa += (dofvals[j] * Np[j])
                end
                E += (JxW) * (pa - pt)^2
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(Ph)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(Ph, qargs)
    return integrate!(elit, qpit, truep)
end

# The function `evaluate_velocity_error` evaluates the true ``L^2``  error of
# the velocity. It does that by integrating the square of the difference
# between the approximate pressure and the  true velocity, the true velocity
# being provided by the `trueux`, `trueuy` functions.
function evaluate_velocity_error(Uxh, Uyh, trueux, trueuy)
    function integrate!(elits, qpits, trueux, trueuy)
        n_du, n_du = ndofsperel.(elits)
        E = 0.0
        for el in zip(elits...)
            uxel, uyel = el
            uxdofvals = eldofvals(uxel)
            uydofvals = eldofvals(uyel)
            for qp in zip(qpits...)
                uxqp, uyqp = qp
                Jac, J = jacjac(uxel, uxqp)
                JxW = J * weight(uxqp)
                Nu = bfun(uxqp)
                uxt = trueux(location(uxel, uxqp)...)
                uyt = trueuy(location(uxel, uxqp)...)
                uxa = 0.0
                uya = 0.0
                for j in 1:n_du
                    uxa += (uxdofvals[j] * Nu[j])
                    uya += (uydofvals[j] * Nu[j])
                end
                E += (JxW) * ((uxa - uxt)^2 + (uya - uyt)^2)
            end
        end
        return sqrt(E)
    end

    elits = (FEIterator(Uxh), FEIterator(Uyh),)
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(Uxh, qargs), QPIterator(Uyh, qargs),)
    return integrate!(elits, qpits, trueux, trueuy)
end

end

# To run the example, evaluate this file which will  compile the module
# `.tut_stokes_ht_p2_p1_reddy`.
using .tut_stokes_ht_p2_p1_reddy
tut_stokes_ht_p2_p1_reddy.run()


text:
abstract type Link end

#Logit link
struct LogitLink <: Link end

function invlink(x::Float64, link::LogitLink)
  logistic(x)
end

function invlink(ax::Array, link::LogitLink)
  logistic.(ax)
end

#Log link
struct LogLink <: Link end

function invlink(x::Float64, link::LogLink)
  exp(x)
end

function invlink(x::Array, link::LogLink)
  exp.(x)
end

ext: jl

max_stars_repo_name: kenkellner/julia-hme

max_issues_repo_name: kenkellner/julia-hme

max_forks_repo_name: kenkellner/julia-hme


text:
using FractalTools
using Makie

# Contruct a tesslation
tridln = Node(TriDelaunay([0., 0.], [1., 0], [0.5, 1.], addboundarypoints=false))

# Definea an observables
msh = lift(tridln) do val
    @info "change in tridln"
    tomesh(val.delaunay)
end
vtx = lift(msh) do val
    @info "change in msh"
    val.position
end

# Plot tesselation
scene = Scene(resolution=(500, 500), scale_plot = false)
mesh!(scene, msh, color=:lightblue)
wireframe!(scene, msh, color=:black)
scatter!(scene, vtx, color=:black)

# Add mouseposition interaction
on(scene.events.mouseposition) do pos
    val = to_world(scene, Point(pos))
    @show val
end

# Add mouse button interaction
on(scene.events.mousebuttons) do buttons
   if ispressed(scene, Mouse.left)
        pos = to_world(scene, Point(scene.events.mouseposition[]))
        @info "$pos added"
        addpoint!(tridln[], pos)
        tridln[] = tridln[]   # Trigger change
        push!(vtx[], pos)
   end
end

# Display scene
display(scene)


text:
module PowerSeries

using Compat

import Base: convert, sqrt, exp, log, sin, cos, tan, asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh, csc, sec, cot, acsc, asec, acot, csch, sech, coth, acsch, asech, acoth, gamma, polygamma, floor, ceil, round, sign, abs, <, /, *, ^, +, -

abstract AbstractSeries{T<:Real, N} <: Number

include("generate_type.jl")

series_types = DataType[]

function generate(order)
  while order > length(series_types)
    push!(series_types, generate_type(length(series_types) + 1))
  end
end

constant(x::Real) = x
polyder(x::Real) = zero(typeof(x))

function /{T, S, N}(x::AbstractSeries{T, N}, y::AbstractSeries{S, N})
  ry = restrict(y)
  constant(x)/constant(y) + polyint(polyder(x)/ry - restrict(x)*polyder(y)/(ry*ry))
end

function /(c::Real, y::AbstractSeries)
  ry = restrict(y)
  c/constant(y) - polyint(c*polyder(y)/(ry*ry))
end

function ^{T, S, N}(x::AbstractSeries{T, N}, y::AbstractSeries{S, N})
  rx = restrict(x)
  ry = restrict(y)
  constant(x)^constant(y) + polyint(polyder(x)*rx^(ry - 1)*ry + log(rx)*rx^ry*polyder(y))
end

function _series_pow_const(x, y)
  if y == 0
    return one(x)
  else
    return constant(x)^y + polyint(polyder(x)*restrict(x)^(y - 1)*y)
  end
end

# First two are to fix redundancy warnings
^(x::AbstractSeries, y::Rational) = _series_pow_const(x, y)
^(x::AbstractSeries, y::Integer) = _series_pow_const(x, y)
^(x::AbstractSeries, y::Real) = _series_pow_const(x, y)

^(::Irrational{:e}, y::AbstractSeries) = exp(y)
^(x::Real, y::AbstractSeries) = x^constant(y) + polyint(log(x)*x^restrict(y)*polyder(y))

sqrt(x::AbstractSeries) = sqrt(constant(x)) + polyint(polyder(x)/(2*sqrt(restrict(x))))
exp(x::AbstractSeries) = exp(constant(x)) + polyint(polyder(x)*exp(restrict(x)))
log(x::AbstractSeries) = log(constant(x)) + polyint(polyder(x)/restrict(x))
sin(x::AbstractSeries) = sin(constant(x)) + polyint(polyder(x)*cos(restrict(x)))
cos(x::AbstractSeries) = cos(constant(x)) - polyint(polyder(x)*sin(restrict(x)))
tan(x::AbstractSeries) = tan(constant(x)) + polyint(polyder(x)*sec(restrict(x))^2)
asin(x::AbstractSeries) = asin(constant(x)) + polyint(polyder(x)/sqrt(1 - restrict(x)^2))
acos(x::AbstractSeries) = acos(constant(x)) - polyint(polyder(x)/sqrt(1 - restrict(x)^2))
atan(x::AbstractSeries) = atan(constant(x)) + polyint(polyder(x)/(1 + restrict(x)^2))
sinh(x::AbstractSeries) = sinh(constant(x)) + polyint(polyder(x)*cosh(restrict(x)))
cosh(x::AbstractSeries) = cosh(constant(x)) + polyint(polyder(x)*sinh(restrict(x)))
tanh(x::AbstractSeries) = tanh(constant(x)) + polyint(polyder(x)*sech(restrict(x))^2)
asinh(x::AbstractSeries) = asinh(constant(x)) + polyint(polyder(x)/sqrt(restrict(x)^2 + 1))
acosh(x::AbstractSeries) = acosh(constant(x)) + polyint(polyder(x)/sqrt(restrict(x)^2 - 1))
atanh(x::AbstractSeries) = atanh(constant(x)) + polyint(polyder(x)/(1 - restrict(x)^2))
csc(x::AbstractSeries) = csc(constant(x)) - polyint(polyder(x)*cot(restrict(x))*csc(restrict(x)))
sec(x::AbstractSeries) = sec(constant(x)) + polyint(polyder(x)*tan(restrict(x))*sec(restrict(x)))
cot(x::AbstractSeries) = cot(constant(x)) - polyint(polyder(x)*csc(restrict(x))^2)

acsc(x::AbstractSeries) =
  acsc(constant(x)) - polyint(polyder(x)/(restrict(x)^2*sqrt(1-1/restrict(x)^2)))

asec(x::AbstractSeries) =
  asec(constant(x)) + polyint(polyder(x)/(restrict(x)^2*sqrt(1-1/restrict(x)^2)))

acot(x::AbstractSeries) = acot(constant(x)) - polyint(polyder(x)/(1 + restrict(x)^2))

csch(x::AbstractSeries) =
  csch(constant(x)) - polyint(polyder(x)*coth(restrict(x))*csch(restrict(x)))

sech(x::AbstractSeries) =
  sech(constant(x)) - polyint(polyder(x)*tanh(restrict(x))*sech(restrict(x)))

coth(x::AbstractSeries) = coth(constant(x)) - polyint(polyder(x)*csch(restrict(x))^2)

acsch(x::AbstractSeries) =
  acsch(constant(x)) - polyint(polyder(x)/(restrict(x)^2*sqrt(1+1/restrict(x)^2)))

function asech(x::AbstractSeries)
  rx = restrict(x)
  asech(constant(x)) + polyint(polyder(x)*sqrt((1 - rx)/(1 + rx))/(rx*(rx - 1)))
end

acoth(x::AbstractSeries) = acoth(constant(x)) + polyint(polyder(x)/(1 - restrict(x)^2))

gamma(x::AbstractSeries) =
  gamma(constant(x)) + polyint(polyder(x)*gamma(restrict(x))*polygamma(0, restrict(x)))

_polygamma(k, x) = polygamma(k, constant(x)) + polyint(polyder(x)*polygamma(k + 1, restrict(x)))

# First two defs are to disambiguate with definitions in base
polygamma(k::Integer, x::AbstractSeries) = _polygamma(k, x)
polygamma{T<:Number}(a::AbstractArray{T}, x::AbstractSeries) =
  reshape([_polygamma(a[i], x) for i=1:length(a)], size(a))
polygamma(k, x::AbstractSeries) = _polygamma(k, x)

# TODO, what about the jump points
floor(x::AbstractSeries) = floor(constant(x)) + polyint(polyder(x)*0)
ceil(x::AbstractSeries) = ceil(constant(x)) + polyint(polyder(x)*0)
round(x::AbstractSeries) = round(constant(x)) + polyint(polyder(x)*0)
sign(x::AbstractSeries) = sign(constant(x)) + polyint(polyder(x)*0)
abs(x::AbstractSeries) = abs(constant(x)) + polyint(polyder(x)*sign(restrict(x)))

<(x::AbstractSeries, c::Real) = constant(x) < c
<(c::Real, x::AbstractSeries) = c < constant(x)
<(x::AbstractSeries, y::AbstractSeries) = constant(x) < constant(y)

generate(7)

export series, restrict, constant, polyint, polyval, polyder

end


text:
"""
    OneHotVector([itr])

Constuct a `OneHotVector`, an `AbstractVector{Bool}` containing exactly one
`true`, and all other values `false`.
The vector is stored in memory as two integers and is immutable.
"""
struct OneHotVector <: AbstractVector{Bool}
    len::Int
    index::Int

    OneHotVector(len::Int, index::Int, ::Unsafe) = new(len, index)
end

function OneHotVector(len::Integer, index::Integer)
    lenT = convert(Int, len)
    indT = convert(Int, index)
    lenT < 1 && throw(ArgumentError("Length must be at least one"))
    in(indT, Base.OneTo(lenT)) || throw(ArgumentError("Index must be in 1:len"))
    return OneHotVector(lenT, indT, unsafe)
end

@noinline notrue() = throw(ArgumentError("No trues in vector"))
@noinline twotrue() = throw(ArgumentError("More than one true in argument"))

function OneHotVector(v)
    index = 0
    length = 0
    for elem in v
        length += 1
        b = convert(Bool, elem)
        b & (index != 0) && twotrue()
        b && (index = length)
    end
    index == 0 && notrue()
    return OneHotVector(length, index, unsafe)
end

function Base.hash(x::OneHotVector, h::UInt)
    base = 0x2c6ed6da44a96001 % UInt
    return hash(x.index, hash(x.len, base ⊻ h))
end

Base.size(v::OneHotVector) = (v.len,)
Base.length(v::OneHotVector) = v.len

function Base.getindex(v::OneHotVector, i::Int)
    @boundscheck checkbounds(v, i)
    return i == v.index
end

Base.getindex(v::OneHotVector, ::Colon) = v
Base.argmax(v::OneHotVector) = v.index
Base.argmin(v::OneHotVector) = 1 + ((v.index == 1) & (length(v) != 1))
Base.count(::typeof(identity), v::OneHotVector) = 1
Base.allunique(v::OneHotVector) = length(v) < 3

for f in (:(Base.:+), :(Base.:-))
    @eval function ($f)(v1::AbstractVector, v::OneHotVector)
        promote_shape(v, v1)
        c = copy(v1)
        @inbounds val = c[v.index]
        val = $(f)(val, one(eltype(v1)))
        @inbounds c[v.index] = val
        return c
    end
end

Base.:+(v::OneHotVector, v1::AbstractVector) = v1 + v

function Base.:-(v::OneHotVector, v1::AbstractVector)
    promote_shape(v, v1)
    c = similar(v1)
    @inbounds for i in eachindex(v1)
        c[i] = -v1[i]
    end
    @inbounds c[v.index] += one(eltype(v1))
    return c
end

Base.reverse(v::OneHotVector) = OneHotVector(length(v), length(v)-v.index+1, unsafe)

function Base.findfirst(f::Function, v::OneHotVector)
    ft::Bool = f(true)
    ff::Bool = f(false)
    ft & ff && return 1
    !(ft | ff) && return nothing
    ft && return v.index
    ff && return ifelse(length(v) == 1, nothing, 1 + (v.index == 1))
end

function Base.circshift(v::OneHotVector, s::Integer)
    newindex = Int(mod1(v.index + s, length(v)))
    return OneHotVector(length(v), newindex, unsafe)
end


text:
# To run this script, `cd` to the `./test/fixtures` directory and then, from the Julia terminal, `include("./runner.jl")`.

import JSON

exp = linspace( 0, 308, 5001 )
len = length( exp )

x = Array( Float64, len )
y = rand( len ) - 0.5;
z = Array( Float64, len )
for i in eachindex(x)
	x[i] = (rand()-0.5) * 10^exp[i]
	z[i] = flipsign( x[i], y[i] )
end

data = Dict([
	("x", x),
	("y", y),
	("expected", z)
])

outfile = open( "data.json", "w" )
write( outfile, JSON.json(data) )
close( outfile );

max_stars_repo_name: compute-io/flipsign

max_issues_repo_name: compute-io/flipsign


text:
@testset "WeakRefStrings" begin
    # fix #198, #199
    db = table(WeakRefStrings.StringVector(["A", "B", "A", "B"]), Dagger.distribute(ones(4), 2), names = [:x, :y])
    z = groupreduce(+, db, :x, select = :y)
    @test fieldtype(typeof(z).parameters[1], 1) == String
    @test fieldtype(typeof(z).parameters[2], 1) == String
    @test fieldtype(typeof(z.chunks[1].domain.second).parameters[1], 1) == String
end


text:
### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 12cc2940-0403-11eb-19a7-bb570de58f6f
begin
	using Pkg
	Pkg.activate(mktempdir())
end

# ╔═╡ 15187690-0403-11eb-2dfd-fd924faa3513
begin
	Pkg.add(["Plots", "PlutoUI",])

	using Plots
	plotly()
	using PlutoUI
end

# ╔═╡ 01341648-0403-11eb-2212-db450c299f35
md"_homework 4, version 1_"

# ╔═╡ 06f30b2a-0403-11eb-0f05-8badebe1011d
md"""

# **Homework 4**: _Epidemic modeling I_
`18.S191`, fall 2020

This notebook contains _built-in, live answer checks_! In some exercises you will see a coloured box, which runs a test case on your code, and provides feedback based on the result. Simply edit the code, run it, and the check runs again.

_For MIT students:_ there will also be some additional (secret) test cases that will be run as part of the grading process, and we will look at your notebook and write comments.

Feel free to ask questions!
"""

# ╔═╡ 095cbf46-0403-11eb-0c37-35de9562cebc
# edit the code below to set your name and kerberos ID (i.e. email without @mit.edu)

student = (name = "Jazzy Doe", kerberos_id = "jazz")

# you might need to wait until all other cells in this notebook have completed running.
# scroll around the page to see what's up

# ╔═╡ 03a85970-0403-11eb-334a-812b59c0905b
md"""

Submission by: **_$(student.name)_** ($(student.kerberos_id)@mit.edu)
"""

# ╔═╡ 107e65a4-0403-11eb-0c14-37d8d828b469
md"_Let's create a package environment:_"

# ╔═╡ 1d3356c4-0403-11eb-0f48-01b5eb14a585
html"""

"""

# ╔═╡ df8547b4-0400-11eb-07c6-fb370b61c2b6
md"""
## **Exercise 1:** _Modelling recovery_

In this exercise we will investigate a simple stochastic (probabilistic) model of recovery from an infection and
the time $\tau$ needed to recover. Although this model can be easily studied analytically using probability theory, we will instead use computational methods. (If you know about this distribution already, try to ignore what you know about it!)

In this model, an individual who is infected has a constant probability $p$ to recover each day. If they recover on day $n$ then $\tau$ takes the value $n$. Each time we run a new experiment $\tau$ will take on different values, so $\tau$ is a (discrete) random variable. We thus need to study statistical properties of $\tau$, such as its mean and its probability distribution.

#### Exercise 1.1 - _Probability distributions_

👉 Define the function `bernoulli(p)`, which returns `true` with probability $p$ and `false` with probability $(1 - p)$.

"""

# ╔═╡ 02b0c2fc-0415-11eb-2b40-7bca8ea4eef9
function bernoulli(p::Number)
	rand() < p
end

# ╔═╡ 76d117d4-0403-11eb-05d2-c5ea47d06f43
md"""
👉 Write a function `recovery_time(p)` that returns the time taken until the person recovers.
"""

# ╔═╡ d57c6a5a-041b-11eb-3ab4-774a2d45a891
function recovery_time(p)
	if p ≤ 0
		throw(ArgumentError("p must be positive: p = 0 cannot result in a recovery"))
	end

	# Your code here. See the comment below about the p ≤ 0 case.
	τ = 1
	while !bernoulli(p)
		τ += 1
	end
	return τ
end

# ╔═╡ 4ad79def-f52b-4b6b-88fa-581a976fc605
begin
	samples = [recovery_time(0.2) for _ in 1:256]
	extrema(samples)
end

# ╔═╡ 6db6c894-0415-11eb-305a-c75b119d89e9
md"""
We should always be aware of special cases (sometimes called "boundary conditions"). Make sure *not* to run the code with $p=0$! What would happen in that case? Your code should check for this and throw an `ArgumentError` as follows:

```julia
throw(ArgumentError("..."))
```

with a suitable error message.

"""

# ╔═╡ 6de37d6c-0415-11eb-1b05-85ac820016c7
md"""
👉 What happens for $p=1$?
"""

# ╔═╡ 73047bba-0416-11eb-1047-23e9c3dbde05
interpretation_of_p_equals_one = md"""
The patient is certain to recover on the first day, so τ is 1 .
"""

# ╔═╡ 76f62d64-0403-11eb-27e2-3de58366b619
md"""
#### Exercise 1.2
👉 Write a function `do_experiment(p, N)` that runs the function `recovery_time` `N` times and collects the results into a vector.
"""

# ╔═╡ c5c7cb86-041b-11eb-3360-45463105f3c9
function do_experiment(p, N)
	return map(recovery_time, Iterators.repeated(p, N))
end

# ╔═╡ d8abd2f6-0416-11eb-1c2a-f9157d9760a7
small_experiment = do_experiment(0.5, 20)

# ╔═╡ 771c8f0c-0403-11eb-097e-ab24d0714ad5
md"""
#### Exercise 1.3
👉 Write a function `frequencies(data)` that calculates and returns the frequencies (i.e. probability distribution) of input data.

The input will be an array of integers, **with duplicates**, and the result will be a dictionary that maps each occured value to its frequency in the data.

For example,
```julia
frequencies([7, 8, 9, 7])
```
should give
```julia
Dict(
	7 => 0.5,
	8 => 0.25,
	9 => 0.25
)
```

As with any probability distribution, it should be normalised to $1$, in the sense that the *total* probability should be $1$.
"""

# ╔═╡ 105d347e-041c-11eb-2fc8-1d9e5eda2be0
function frequencies(values)
	n = length(values)
	freq = Dict{Int, Float64}()
	for k in values
		if haskey(freq, k)
			freq[k] += 1/n
		else
			freq[k] = 1/n
		end
	end
	return freq
end

# ╔═╡ 1ca7a8c2-041a-11eb-146a-15b8cdeaea72
frequencies(small_experiment)

# ╔═╡ 77428072-0403-11eb-0068-81e3728f2ebe
md"""
Let's run an experiment with $p=0.25$ and $N=10,000$.
"""

# ╔═╡ 4b3ec86c-0419-11eb-26fd-cbbfdf19afa8
large_experiment = do_experiment(0.25, 10_000)
# (10_000 is just 10000 but easier to read)

# ╔═╡ dc784864-0430-11eb-1478-d1153e017310
md"""
The frequencies dictionary is difficult to interpret on its own, so instead, we will **plot** it, i.e. plot $P(\tau = n)$ against $n$, where $n$ is the recovery time.

Plots.jl comes with a function `bar`, which does exactly what we want:
"""

# ╔═╡ 8a28c56e-04b4-11eb-279c-3b4dfb2a9f9b
bar(frequencies(large_experiment))

# ╔═╡ 9374e63c-0493-11eb-0952-4b97512d7cdb
md"""
Great! Feel free to experiment with this function, try giving it a different array as argument. Plots.jl is pretty clever, it even works with an array of strings!

#### Exercise 1.4
Next, we want to **add a new element** to our plot: a vertical line. To demonstrate how this works, here we added a vertical line at the _maximum value_.

To write this function, we first create a **base plot**, we then **modify** that plot to add the vertical line, and finally, we **return** the plot. More on this in [the next info box](#note_about_plotting).
"""

# ╔═╡ 823364ce-041c-11eb-2467-7ffa4f751527
function frequencies_plot_with_maximum(data::Vector)
	base = bar(frequencies(data))
	vline!(base, [maximum(data)], label="maximum")

	return base
end

# ╔═╡ 1ddbaa18-0494-11eb-1fc8-250ab6ae89f1
frequencies_plot_with_maximum(large_experiment)

# ╔═╡ f3f81172-041c-11eb-2b9b-e99b7b9400ed
md"""
$(html"")
> ### Note about plotting
>
> Plots.jl has an interesting property: a plot is an object, not an action. Functions like `plot`, `bar`, `histogram` don't draw anything on your screen - they just return a `Plots.Plot`. This is a struct that contains the _description_ of a plot (what data should be plotted in what way?), not the _picture_.
>
> So a Pluto cell with a single line, `plot(1:10)`, will show a plot, because the _result_ of the function `plot` is a `Plot` object, and Pluto just shows the result of a cell.
>
> ##### Modifying plots
> Nice plots are often formed by overlaying multiple plots. In Plots.jl, this is done using the **modifying functions**: `plot!`, `bar!`, `vline!`, etc. These take an extra (first) argument: a previous plot to modify.
>
> For example, to plot the `sin`, `cos` and `tan` functions in the same view, we do:
> ```julia
> function sin_cos_plot()
>     T = -1.0:0.01:1.0
>
>     result = plot(T, sin.(T))
>     plot!(result, T, cos.(T))
>     plot!(result, T, tan.(T))
>
>     return result
> end
> ```
>
> 💡 This example demonstrates a useful pattern to combine plots:
> 1. Create a **new** plot and store it in a variable
> 2. **Modify** that plot to add more elements
> 3. Return the plot
>
> It is recommended that these 3 steps happen **within a single cell**. This can prevent some strange glitches when re-running cells. There are three ways to group expressions together into a single cell: `begin`, `let` and `function`. More on this [later](#function_begin_let)!
"""

# ╔═╡ 7768a2dc-0403-11eb-39b7-fd660dc952fe
md"""
👉 Write the function `frequencies_plot_with_mean` that calculates the mean recovery time and displays it using a vertical line.
"""

# ╔═╡ f1f89502-0494-11eb-2303-0b79d8bbd13f
function frequencies_plot_with_mean(data)
	# start out by copying the frequencies_plot_with_maximum function
	base = bar(frequencies(data), label="Relative Frequency")
	vline!(base, [sum(data)/length(data)], label="Mean")
	return base
end

# ╔═╡ 06089d1e-0495-11eb-0ace-a7a7dc60e5b2
frequencies_plot_with_mean(large_experiment)

# ╔═╡ 77b54c10-0403-11eb-16ad-65374d29a817
md"""
👉 Write an interactive visualization that draws the histogram and mean for $p$ between $0.01$ (not $0$!) and $1$, and $N$ between $1$ and $100,000$, say. To avoid a naming conflict, call them `p_interactive` and `N_interactive`, instead of just `p` and `N`.
"""

# ╔═╡ 103ed025-823f-4923-8e5f-55b2b694fe70
@bind p_interactive Slider(0.01:0.01:1, show_value=true)

# ╔═╡ f33f2b39-ed8a-4808-9422-1bd94718e76c
@bind N_interactive Slider(1:100:100000, show_value=true)

# ╔═╡ bb63f3cc-042f-11eb-04ff-a128aec3c378
begin
	exp_data = do_experiment(p_interactive, N_interactive)
	frequencies_plot_with_mean(exp_data)
end

# ╔═╡ bb8aeb58-042f-11eb-18b8-f995631df619
md"""
As you separately vary $p$ and $N$, what do you observe about the **mean** in each case? Does that make sense?
"""

# ╔═╡ 992aadc6-4e4e-4f9c-9091-557d2a54bd25
md"""As p varies the mean decreases tending towards 1. This makes sense because the expected value of the time to first success of a Bernoulli process is 1/p. As N increases the mean decreases and tends towards the true value of the mean. This is a realisation of the law of large numbers.
"""

# ╔═╡ 778ec25c-0403-11eb-3146-1d11c294bb1f
md"""
#### Exercise 1.5
👉 What shape does the distribution seem to have? Can you verify that by adding a second plot with the expected shape?
"""

# ╔═╡ 5c9c0050-9f2a-41b6-9c71-b6fa4658330c
md"""The distribution has an exponential decay shape of the form $A \mathrm{e}^{kx}$ We can apply  least-squares regression to the log of the distribution to obtain an estimate of $A$ and $k$.
"""

# ╔═╡ f1693d9c-a6a7-4e0d-b958-1a6a459882d7
function fit_exponential(x, y)
	w = log.(y)
	n_x = length(x)
	b = (sum(w .* x) - (sum(w) * sum(x))/n_x) / (sum(x.^2) - sum(x)^2 / n_x)
	a = sum(w) / n_x - (b / n_x * sum(x))
	return (exp(a), b)
end

# ╔═╡ b8a2f491-fc02-4ad1-a762-7d36643ce844
begin
	sorted_frequencies = sort(frequencies(exp_data))
	xs = keys(sorted_frequencies) |> collect
	ys = values(sorted_frequencies)
	A, k = fit_exponential(xs, ys)
	bar_plot = frequencies_plot_with_mean(exp_data)
	#xs = pushfirst!(xs, 0)
	plot!(bar_plot, xs, A*exp.(k*xs), label="Fit")
end

# ╔═╡ 7bb8e426-0495-11eb-3a8b-cbbab61a1631
sort(frequencies(exp_data))

# ╔═╡ 77db111e-0403-11eb-2dea-4b42ceed65d6
md"""
#### Exercise 1.6
👉 Use $N = 10,000$ to calculate the mean time $\langle \tau(p) \rangle$ to recover as a function of $p$ between $0.001$ and $1$ (say). Plot this relationship.

"""

# ╔═╡ 321ac3d6-b40a-49d0-bbb5-5f1a9f6cd6dc
function my_mean(x)
	sum(x) / length(x)
end

# ╔═╡ 3ea098eb-3dd2-4e1e-be87-8aad8ecc5dfb
begin
	ps = 0.001:0.001:1
	mean_time = map(ps) do p
		data = do_experiment(p, 10_000)
		my_mean(data)
	end
	plot(ps, mean_time, title="Mean time to recover", xlabel="Probability", ylabel="Time / day")
end

# ╔═╡ 61789646-0403-11eb-0042-f3b8308f11ba
md"""
## **Exercise 2:** _Agent-based model for an epidemic outbreak -- types_

In this and the following exercises we will develop a simple stochastic model for combined infection and recovery in a population, which may exhibit an **epidemic outbreak** (i.e. a large spike in the number of infectious people).
The population is **well mixed**, i.e. everyone is in contact with everyone else.
[An example of this would be a small school or university in which people are
constantly moving around and interacting with each other.]

The model is an **individual-based** or **agent-based** model:
we explicitly keep track of each individual, or **agent**, in the population and their
infection status. For the moment we will not keep track of their position in space;
we will just assume that there is some mechanism, not included in the model, by which they interact with other individuals.

#### Exercise 2.1

Each agent will have its own **internal state**, modelling its infection status, namely "susceptible", "infectious" or "recovered". We would like to code these as values `S`, `I` and `R`, respectively. One way to do this is using an [**enumerated type**](https://en.wikipedia.org/wiki/Enumerated_type) or **enum**. Variables of this type can take only a pre-defined set of values; the Julia syntax is as follows:
"""

# ╔═╡ 26f84600-041d-11eb-1856-b12a3e5c1dc7
@enum InfectionStatus S I R

# ╔═╡ 271ec5f0-041d-11eb-041b-db46ec1465e0
md"""
We have just defined a new type `InfectionStatus`, as well as names `S`, `I` and `R` that are the (only) possible values that a variable of this type can take.

👉 Define a variable `test_status` whose value is `S`.
"""

# ╔═╡ 7f4e121c-041d-11eb-0dff-cd0cbfdfd606
test_status = S

# ╔═╡ 7f744644-041d-11eb-08a0-3719cc0adeb7
md"""
👉 Use the `typeof` function to find the type of `test_status`.
"""

# ╔═╡ 88c53208-041d-11eb-3b1e-31b57ba99f05
typeof(test_status)

# ╔═╡ 847d0fc2-041d-11eb-2864-79066e223b45
md"""
👉 Convert `x` to an integer using the `Integer` function. What value does it have? What values do `I` and `R` have?
"""

# ╔═╡ cc8a6234-27e0-49c7-af1a-881d98a85ec6
Integer(test_status)

# ╔═╡ 860790fc-0403-11eb-2f2e-355f77dcc7af
md"""
#### Exercise 2.2

For each agent we want to keep track of its infection status and the number of *other* agents that it infects during the simulation. A good solution for this is to define a *new type* `Agent` to hold all of the information for one agent, as follows:
"""

# ╔═╡ ae4ac4b4-041f-11eb-14f5-1bcde35d18f2
begin
	mutable struct Agent
		status::InfectionStatus
		num_infected::Int64
	end
	Agent() = Agent(S, 0)
end

# ╔═╡ ae70625a-041f-11eb-3082-0753419d6d57
md"""
When you define a new type like this, Julia automatically defines one or more **constructors**, which are methods of a generic function with the *same name* as the type. These are used to create objects of that type.

👉 Use the `methods` function to check how many constructors are pre-defined for the `Agent` type.
"""

# ╔═╡ 60a8b708-04c8-11eb-37b1-3daec644ac90
methods(Agent)

# ╔═╡ 189cae1e-0424-11eb-2666-65bf297d8bdd
md"""
👉 Create an agent `test_agent` with status `S` and `num_infected` equal to 0.
"""

# ╔═╡ 18d308c4-0424-11eb-176d-49feec6889cf
test_agent = S

# ╔═╡ 190deebc-0424-11eb-19fe-615997093e14
md"""
👉 For convenience, define a new constructor (i.e. a new method for the function) that takes no arguments and creates an `Agent` with status `S` and number infected 0, by calling one of the default constructors that Julia creates. This new method lives *outside* (not inside) the definition of the `struct`. (It is called an **outer constructor**.)

(In Pluto, multiple methods for the same function need to be combined in a single cell using a `begin end` block.)

Let's check that the new method works correctly. How many methods does the constructor have now?

"""

# ╔═╡ b329070b-f735-47c7-84fa-d0f60dbdd666
md"""It has 3 methods"""

# ╔═╡ 82f2580a-04c8-11eb-1eea-bdb4e50eee3b
Agent()

# ╔═╡ 8631a536-0403-11eb-0379-bb2e56927727
md"""
#### Exercise 2.3
👉 Write functions `set_status!(a)` and `set_num_infected!(a)` which modify the respective fields of an `Agent`. Check that they work. [Note the bang ("`!`") at the end of the function names to signify that these functions *modify* their argument.]

"""

# ╔═╡ 98beb336-0425-11eb-3886-4f8cfd210288
function set_status!(agent::Agent, new_status::InfectionStatus)
	agent.status = new_status
end

# ╔═╡ 26d27f54-bb4f-45e2-9e67-6f54621c3f18
function set_num_infected!(agent::Agent, new_num_infected::Int64)
	agent.num_infected = new_num_infected
end

# ╔═╡ 866299e8-0403-11eb-085d-2b93459cc141
md"""
👉 We will also need functions `is_susceptible` and `is_infected` that check if a given agent is in those respective states.

"""

# ╔═╡ 9a837b52-0425-11eb-231f-a74405ff6e23
function is_susceptible(agent::Agent)

	return agent.status == S
end

# ╔═╡ a8dd5cae-0425-11eb-119c-bfcbf832d695
function is_infected(agent::Agent)

	return agent.status == I
end

# ╔═╡ 8692bf42-0403-11eb-191f-b7d08895274f
md"""
#### Exericse 2.4
👉 Write a function `generate_agents(N)` that returns a vector of `N` freshly created `Agent`s. They should all be initially susceptible, except one, chosen at random (i.e. uniformly), who is infectious.

"""

# ╔═╡ 7946d83a-04a0-11eb-224b-2b315e87bc84
function generate_agents(N::Integer)
	agents = Array{Agent}(undef, N)
	for i in eachindex(agents)
		agents[i] = Agent()
	end
	infectious_i = rand(1:N)
	set_status!(agents[infectious_i], I)
	return agents
end

# ╔═╡ 488771e2-049f-11eb-3b0a-0de260457731
generate_agents(3)

# ╔═╡ 86d98d0a-0403-11eb-215b-c58ad721a90b
md"""
We will also need types representing different infections.

Let's define an (immutable) `struct` called `InfectionRecovery` with parameters `p_infection` and `p_recovery`. We will make it a subtype of an abstract `AbstractInfection` type, because we will define more infection types later.
"""

# ╔═╡ 223933a4-042c-11eb-10d3-852229f25a35
abstract type AbstractInfection end

# ╔═╡ 1a654bdc-0421-11eb-2c38-7d35060e2565
struct InfectionRecovery <: AbstractInfection
	p_infection
	p_recovery
end

# ╔═╡ 2d3bba2a-04a8-11eb-2c40-87794b6aeeac
md"""
#### Exercise 2.5
👉 Write a function `interact!` that takes an affected `agent` of type `Agent`, an `source` of type `Agent` and an `infection` of type `InfectionRecovery`.  It implements a single (one-sided) interaction between two agents:

- If the `agent` is susceptible and the `source` is infectious, then the `source` infects our `agent` with the given infection probability. If the `source` successfully infects the other agent, then its `num_infected` record must be updated.
- If the `agent` is infected then it recovers with the relevant probability.
- Otherwise, nothing happens.

$(html"")
"""

# ╔═╡ b21475c6-04ac-11eb-1366-f3b5e967402d
md"""
Play around with the test case below to test your function! Try changing the definitions of `agent`, `source` and `infection`. Since we are working with randomness, you might want to run the cell multiple times.
"""

# ╔═╡ 619c8a10-0403-11eb-2e89-8b0974fb01d0
md"""
## **Exercise 3:** _Agent-based model for an epidemic outbreak --  Monte Carlo simulation_

In this exercise we will build on Exercise 2 to write a Monte Carlo simulation of how an infection propagates in a population.

Make sure to re-use the functions that we have already written, and introduce new ones if they are helpful! Short functions make it easier to understand what the function does and build up new functionality piece by piece.

You should not use any global variables inside the functions: Each function must accept as arguments all the information it requires to carry out its task. You need to think carefully about what the information each function requires.

#### Exercise 3.1

👉 Write a function `step!` that takes a vector of `Agent`s and an `infection` of type `InfectionRecovery`. It implements a single step of the infection dynamics as follows:

- Choose two random agents: an `agent` and a `source`.
- Apply `interact!(agent, source, infection)`.
- Return `agents`.

"""

# ╔═╡ 955321de-0403-11eb-04ce-fb1670dfbb9e
md"""
👉 Write a function `sweep!`. It runs `step!` $N$ times, where $N$ is the number of agents. Thus each agent acts, on average, once per sweep; a sweep is thus the unit of time in our Monte Carlo simulation.
"""

# ╔═╡ 95771ce2-0403-11eb-3056-f1dc3a8b7ec3
md"""
👉 Write a function `simulation` that does the following:

1. Generate the $N$ agents.

2. Run `sweep!` a number $T$ of times. Calculate and store the total number of agents with each status at each step in variables `S_counts`, `I_counts` and `R_counts`.

3. Return the vectors `S_counts`, `I_counts` and `R_counts` in a **named tuple**, with keys `S`, `I` and `R`.

You've seen an example of named tuples before: the `student` variable at the top of the notebook!

_Feel free to store the counts in a different way, as long as the return type is the same._
"""

# ╔═╡ 28db9d98-04ca-11eb-3606-9fb89fa62f36
@bind run_basic_sir Button("Run simulation again!")

# ╔═╡ 0a967f38-0493-11eb-0624-77e40b24d757
md"""
We used a `let` block in this cell to group multiple expressions together, but how is it different from `begin` or `function`?

$(html"")
> ##### _**function**_ vs. _**begin**_ vs. _**let**_
> Writing functions is a way to group multiple expressions (i.e. lines of code) together into a mini-program. Note the following about functions:
> - A function always returns **one object**.[^1] This object can be given explicitly by writing `return x`, or implicitly: Julia functions always return the result of the last expression by default. So `f(x) = x+2` is the same as `f(x) = return x+2`.
> - Variables defined inside a function are _not accessible outside the function_. We say that function bodies have a **local scope**. This helps to keep your program easy to read and write: if you define a local variable, then you don't need to worry about it in the rest of the notebook.
>
> There are two other ways to group epxressions together that you might have seen before: `begin` and `let`.
>
> ###### begin
> **`begin`** will group expressions together, and it takes the value of its last subexpression.
>
> We use it in this notebook when we want multiple expressions to always run together.
>
> ###### let
> **`let`** also groups multiple expressions together into one, but variables defined inside of it are **local**: they don't affect code outside of the block. So like `begin`, it is just a block of code, but like `function`, it has a local variable scope.
>
> We use it when we want to define some local (temporary) variables to produce a complicated result, without interfering with other cells. Pluto allows only one definition per _global_ variable of the same name, but you can define _local_ variables with the same names whenever you wish!
>
> [^1]: Even a function like
>
>     `f(x) = return`
>
>     returns **one object**: the object `nothing` — try it out!
"""

# ╔═╡ bf6fd176-04cc-11eb-008a-2fb6ff70a9cb
md"""
#### Exercise 3.2
Alright! Every time that we run the simulation, we get slightly different results, because it is based on randomness. By running the simulation a number of times, you start to get an idea of the _mean behaviour_ of our model. This is the essence of a Monte Carlo method! You use computer-generated randomness to generate samples.

Instead of pressing the button many times, let's have the computer repeat the simulation. In the next cells, we run your simulation `num_simulations=20` times with $N=100$, $p_\text{infection} = 0.02$, $p_\text{infection} = 0.002$ and $T = 1000$.

Every single simulation returns a named tuple with the status counts, so the result of multiple simulations will be an array of those. Have a look inside the result, `simulations`, and make sure that its structure is clear.
"""

# ╔═╡ 38b1aa5a-04cf-11eb-11a2-930741fc9076
function repeat_simulations(f, N, T, infection, num_simulations)

	map(1:num_simulations) do _
		f(N, T, infection)
	end
end

# ╔═╡ 80e6f1e0-04b1-11eb-0d4e-475f1d80c2bb
md"""
In the cell below, we plot the evolution of the number of $I$ individuals as a function of time for each of the simulations on the same plot using transparency (`alpha=0.5` inside the plot command).
"""

# ╔═╡ 95c598d4-0403-11eb-2328-0175ed564915
md"""
👉 Write a function `sir_mean_plot` that returns a plot of the means of $S$, $I$ and $R$ as a function of time on a single graph.
"""

# ╔═╡ 843fd63c-04d0-11eb-0113-c58d346179d6
function sir_mean_plot(simulations::Vector{<:NamedTuple})
	# you might need T for this function, here's a trick to get it:
	T = length(first(simulations).S)
	p = plot()

	n_sims = length(simulations)
	for i_status in eachindex(first(simulations))
		mean_i = sum(s -> getfield(s, i_status), simulations) ./ n_sims
		plot!(p, 1:T, mean_i, lw=3, alpha=1, label="Mean type $(string(i_status))")
	end
	return p
end

# ╔═╡ dfb99ace-04cf-11eb-0739-7d694c837d59
md"""
👉 Allow $p_\text{infection}$ and $p_\text{recovery}$ to be changed interactively and find parameter values for which you observe an epidemic outbreak.
"""

# ╔═╡ 1c6aa208-04d1-11eb-0b87-cf429e6ff6d0
@bind p_infection Slider(0:0.01:0.5, default=0.08, show_value=true)

# ╔═╡ 19f75d2b-e35c-4475-a7b2-bc9d224075dc
@bind p_recovery Slider(0:0.001:0.2, default=0.004, show_value=true)

# ╔═╡ 95eb9f88-0403-11eb-155b-7b2d3a07cff0
md"""
👉 Write a function `sir_mean_error_plot` that does the same as `sir_mean_plot`, which also computes the **standard deviation** $\sigma$ of $S$, $I$, $R$ at each step. Add this to the plot using **error bars**, using the option `yerr=σ` in the plot command; use transparency.

This should confirm that the distribution of $I$ at each step is pretty wide!
"""

# ╔═╡ 287ee7aa-0435-11eb-0ca3-951dbbe69404
function sir_mean_error_plot(simulations::Vector{<:NamedTuple})
	# you might need T for this function, here's a trick to get it:
	T = length(first(simulations).S)
	n_sims = length(simulations)
	# define helper functions
	fetch_means(i_status) = sum(getfield.(simulations, i_status)) ./ n_sims
	my_std(v) = sqrt(sum((v .- my_mean(v)).^2) / length(v))
	function fetch_std(i_status)
		sims_i_status = getfield.(simulations, i_status)
		σ = my_std.([(s[i] for s in sims_i_status) for i in 1:T])
		return σ
	end
	p = plot()
	for i_status in eachindex(first(simulations))
		mean_i = fetch_means(i_status)
		σ = fetch_std(i_status)
		plot!(p, 1:T, mean_i, ribbon=σ, lw=2, alpha=1, fillalpha=0.3, label="Mean type $(string(i_status))")
	end
	return p
end

# ╔═╡ 9611ca24-0403-11eb-3582-b7e3bb243e62
md"""
#### Exercise 3.3

👉 Plot the probability distribution of `num_infected`. Does it have a recognisable shape? (Feel free to increase the number of agents in order to get better statistics.)

"""

# ╔═╡ f9cff122-d647-4a3f-861a-4a5967922368
md"""The distribution has an exponential decay shape."""

# ╔═╡ 9635c944-0403-11eb-3982-4df509f6a556
md"""
#### Exercse 3.4
👉 What are three *simple* ways in which you could characterise the magnitude (size) of the epidemic outbreak? Find approximate values of these quantities for one of the runs of your simulation.

"""

# ╔═╡ 4ad11052-042c-11eb-3643-8b2b3e1269bc
md"""
1. The mean number of agents a single infected agent infects
2. The mean time to recover
3. The peak number of infected agents
"""

# ╔═╡ 61c00724-0403-11eb-228d-17c11670e5d1
md"""
## **Exercise 4:** _Reinfection_

In this exercise we will *re-use* our simulation infrastructure to study the dynamics of a different type of infection: there is no immunity, and hence no "recovery" rather, susceptible individuals may now be **re-infected**

#### Exercise 4.1
👉 Make a new infection type `Reinfection`. This has the *same* two fields as `InfectionRecovery` (`p_infection` and `p_recovery`). However, "recovery" now means "becomes susceptible again", instead of "moves to the `R` class.

This new type `Reinfection` should also be a **subtype** of `AbstractInfection`. This allows us to reuse our previous functions, which are defined for the abstract supertype.
"""

# ╔═╡ 8dd97820-04a5-11eb-36c0-8f92d4b859a8
struct Reinfection <: AbstractInfection
	p_infection::Real
	p_recovery::Real
end

# ╔═╡ 406aabea-04a5-11eb-06b8-312879457c42
begin
	function interact!(agent::Agent, source::Agent, infection::InfectionRecovery)
		if is_susceptible(agent) && is_infected(source)
			if bernoulli(infection.p_infection)
				set_status!(agent, I)
				set_num_infected!(source, source.num_infected+1)
				end
		elseif is_infected(agent)
			if bernoulli(infection.p_recovery)
				set_status!(agent, R)
			end
		end
	end

	function interact!(agent::Agent, source::Agent, infection::Reinfection)
		if is_susceptible(agent) && is_infected(source)
			if bernoulli(infection.p_infection)
				set_status!(agent, I)
				set_num_infected!(source, source.num_infected+1)
				end
		elseif is_infected(agent)
			if bernoulli(infection.p_recovery)
				set_status!(agent, S)
			end
		end
	end
end

# ╔═╡ 9c39974c-04a5-11eb-184d-317eb542452c
let
	agent = Agent(S, 0)
	source = Agent(I, 0)
	infection = InfectionRecovery(0.2, 0.6)

	interact!(agent, source, infection)

	(agent=agent, source=source)
end

# ╔═╡ 2ade2694-0425-11eb-2fb2-390da43d9695
function step!(agents::Vector{Agent}, infection::AbstractInfection)
	agent = rand(agents)
	source = rand(agents)
	interact!(agent, source, infection)
	return agents
end

# ╔═╡ 46133a74-04b1-11eb-0b46-0bc74e564680
function sweep!(agents::Vector{Agent}, infection::AbstractInfection)
	N = length(agents)
	for _ in 1:N
		step!(agents, infection)
	end
end

# ╔═╡ 887d27fc-04bc-11eb-0ab9-eb95ef9607f8
function simulation(N::Integer, T::Integer, infection::AbstractInfection)

	agents = generate_agents(N)
	S_counts, I_counts, R_counts = zeros(Int, T), zeros(Int, T), zeros(Int, T)
	for t in 1:T
		sweep!(agents, infection)
		S_counts[t] = count(a -> is_susceptible(a), agents)
		I_counts[t] = count(a -> is_infected(a), agents)
		R_counts[t] = count(a -> a.status == R, agents)
	end

	return (S=S_counts, I=I_counts, R=R_counts)
end

# ╔═╡ b92f1cec-04ae-11eb-0072-3535d1118494
simulation(3, 20, InfectionRecovery(0.9, 0.2))

# ╔═╡ 2c62b4ae-04b3-11eb-0080-a1035a7e31a2
simulation(100, 1000, InfectionRecovery(0.005, 0.2))

# ╔═╡ c5156c72-04af-11eb-1106-b13969b036ca
let
	run_basic_sir

	N = 100
	T = 1000
	sim = simulation(N, T, InfectionRecovery(0.02, 0.002))

	result = plot(1:T, sim.S, ylim=(0, N), label="Susceptible")
	plot!(result, 1:T, sim.I, ylim=(0, N), label="Infectious")
	plot!(result, 1:T, sim.R, ylim=(0, N), label="Recovered")
end

# ╔═╡ 80c2cd88-04b1-11eb-326e-0120a39405ea
simulations = repeat_simulations(simulation, 100, 1000, InfectionRecovery(p_infection, p_recovery), 20)

# ╔═╡ ca8e82da-6d97-4f02-9c66-8fd63349458b
let
	p = plot()

	for sim in simulations
		plot!(p, 1:1000, sim.I, alpha=.5, label=nothing)
	end

	n_sims = length(simulations)
	mean_I = sum(s -> s.I, simulations) ./ n_sims
	plot!(p, 1:1000, mean_I, lw=3, alpha=1, label="Mean")
	p
end

# ╔═╡ 7f635722-04d0-11eb-3209-4b603c9e843c
sir_mean_plot(simulations)

# ╔═╡ 2efb95c6-0f45-4c6a-a808-a4046daea0f8
sir_mean_error_plot(simulations)

# ╔═╡ 26e2978e-0435-11eb-0d61-25f552d2771e
function num_infected_sim(N::Integer, T::Integer, infection::AbstractInfection)

	agents = generate_agents(N)
	for t in 1:T
		sweep!(agents, infection)
	end
	return getfield.(agents, :num_infected)
end

# ╔═╡ c72e0b01-8730-4b24-aa1a-ba6e54326ad5
frequencies_plot_with_mean(Iterators.flatten(repeat_simulations(num_infected_sim, 5000, 5000, InfectionRecovery(p_infection, p_recovery), 20)) |> collect)

# ╔═╡ 99ef7b2a-0403-11eb-08ef-e1023cd151ae
md"""
👉 Make a *new method* for the `interact!` function that accepts the new infection type as argument, reusing as much functionality as possible from the previous version.

Write it in the same cell as [our previous `interact!` method](#interactfunction), and use a `begin` block to group the two definitions together.

"""

# ╔═╡ 9a13b17c-0403-11eb-024f-9b37e95e211b
md"""
#### Exercise 4.2
👉 Run the simulation 20 times and plot $I$ as a function of time for each one, together with the mean over the 20 simulations (as you did in the previous exercises).

Note that you should be able to re-use the `sweep!` and `simulation` functions , since those should be sufficiently **generic** to work with the new `step!` function! (Modify them if they are not.)

"""

# ╔═╡ 1ac4b33a-0435-11eb-36f8-8f3f81ae7844

# ╔═╡ 9a377b32-0403-11eb-2799-e7e59caa6a45
md"""
👉 Run the new simulation and draw $I$ (averaged over runs) as a function of time. Is the behaviour qualitatively the same or different? Describe what you see.

"""

# ╔═╡ 21c50840-0435-11eb-1307-7138ecde0691

# ╔═╡ da49710e-0420-11eb-092e-4f1173868738
md"""
## **Exercise 5** - _Lecture transcript_
(MIT students only)
Please see the link for hw 4 transcript document on [Canvas](https://canvas.mit.edu/courses/5637).
We want each of you to correct about 400 lines, but don’t spend more than 15 minutes on it.
See the the beginning of the document for more instructions.
:point_right: Please mention the name of the video(s) and the line ranges you edited:
"""

# ╔═╡ e6219c7c-0420-11eb-3faa-13126f7c8007
lines_i_edited = md"""
Abstraction, lines 1-219

Array Basics, lines 1-137

Course Intro, lines 1-44

(_for example_)
"""

# ╔═╡ 531d13c2-0414-11eb-0acd-4905a684869d
if student.name == "Jazzy Doe"
	md"""
	!!! danger "Before you submit"
	    Remember to fill in your **name** and **Kerberos ID** at the top of this notebook.
	"""
end

# ╔═╡ 4f19e872-0414-11eb-0dfd-e53d2aecc4dc
md"## Function library

Just some helper functions used in the notebook."

# ╔═╡ 48a16c42-0414-11eb-0e0c-bf52bbb0f618
hint(text) = Markdown.MD(Markdown.Admonition("hint", "Hint", [text]))

# ╔═╡ 6d906d0c-0415-11eb-0c1c-b5c0aca841db
hint(md"Remember to always re-use work you have done previously: in this case you should re-use the function `bernoulli`.")

# ╔═╡ 08e2bc64-0417-11eb-1457-21c0d18e8c51
hint(md"""
Do you remember how we worked with dictionaries in Homework 3? You can create an empty dictionary using `Dict()`. You may want to use either the function `haskey` or the function `get` on your dictionary -- check the documentation for how to use these functions.
""")

# ╔═╡ 9cf9080a-04b1-11eb-12a0-17013f2d37f5
md"""
👉 Calculate the **mean number of infectious agents** of our simulations for each time step. Add it to the plot using a heavier line (`lw=3` for "linewidth") by modifying the cell above.

Check the answer yourself: does your curve follow the average trend?

$(hint(md"This exercise requires some creative juggling with arrays, anonymous functions, `map`s, or whatever you see fit!"))
"""

# ╔═╡ 461586dc-0414-11eb-00f3-4984b57bfac5
almost(text) = Markdown.MD(Markdown.Admonition("warning", "Almost there!", [text]))

# ╔═╡ 43e6e856-0414-11eb-19ca-07358aa8b667
still_missing(text=md"Replace `missing` with your answer.") = Markdown.MD(Markdown.Admonition("warning", "Here we go!", [text]))

# ╔═╡ 41cefa68-0414-11eb-3bad-6530360d6f68
keep_working(text=md"The answer is not quite right.") = Markdown.MD(Markdown.Admonition("danger", "Keep working on it!", [text]))

# ╔═╡ 3f5e0af8-0414-11eb-34a7-a71e7aaf6443
yays = [md"Fantastic!", md"Splendid!", md"Great!", md"Yay ❤", md"Great! 🎉", md"Well done!", md"Keep it up!", md"Good job!", md"Awesome!", md"You got the right answer!", md"Let's move on to the next section."]

# ╔═╡ 3d88c056-0414-11eb-0025-05d3aff1588b
correct(text=rand(yays)) = Markdown.MD(Markdown.Admonition("correct", "Got it!", [text]))

# ╔═╡ 3c0528a0-0414-11eb-2f68-a5657ab9e73d
not_defined(variable_name) = Markdown.MD(Markdown.Admonition("danger", "Oopsie!", [md"Make sure that you define a variable called **$(Markdown.Code(string(variable_name)))**"]))

# ╔═╡ b817f466-04d4-11eb-0a26-c1c667f9f7f7
if !@isdefined(bernoulli)
	not_defined(:bernoulli)
else
	let
		result = bernoulli(0.5)

		if result isa Missing
			still_missing()
		elseif !(result isa Bool)
			keep_working(md"Make sure that you return either `true` or `false`.")
		else
			if bernoulli(0.0) == false && bernoulli(1.0) == true
				correct()
			else
				keep_working()
			end
		end
	end
end

# ╔═╡ c61f35ea-04d6-11eb-2503-17a79f8d0298
if !@isdefined(recovery_time)
	not_defined(:recovery_time)
else
	let
		result = recovery_time(1.0)

		if result isa Missing
			still_missing()
		elseif !(result isa Integer)
			keep_working(md"Make sure that you return an integer: the recovery time.")
		else
			if result == 1
				samples = [recovery_time(0.2) for _ in 1:256]

				a, b = extrema(samples)
				if a == 1 && b > 20
					correct()
				else
					keep_working()
				end
			else
				keep_working(md"`p = 1.0` should return `1`: the agent recovers after the first time step.")
			end
		end
	end
end

# ╔═╡ 7c515a7a-04d5-11eb-0f36-4fcebff709d5
if !@isdefined(set_status!)
	not_defined(:set_status!)
else
	let
		agent = Agent(I,2)

		set_status!(agent, R)

		if agent.status == R
			correct()
		else
			keep_working()
		end
	end
end

# ╔═╡ c4a8694a-04d4-11eb-1eef-c9e037e6b21f
if !@isdefined(is_susceptible)
	not_defined(:is_susceptible)
else
	let
		result1 = is_susceptible(Agent(I,2))
		result2 = is_infected(Agent(I,2))

		if result1 isa Missing || result2 isa Missing
			still_missing()
		elseif !(result1 isa Bool) || !(result2 isa Bool)
			keep_working(md"Make sure that you return either `true` or `false`.")
		elseif result1 === false && result2 === true
			if is_susceptible(Agent(S,3)) && !is_infected(Agent(R,9))
				correct()
			else
				keep_working()
			end
		else
			keep_working()
		end
	end
end

# ╔═╡ 393041ec-049f-11eb-3089-2faf378445f3
if !@isdefined(generate_agents)
	not_defined(:generate_agents)
else
	let
		result = generate_agents(4)

		if result isa Missing
			still_missing()
		elseif result isa Nothing
			keep_working("The function returned `nothing`. Did you forget to return something?")
		elseif !(result isa Vector) || !all(x -> x isa Agent, result)
			keep_working(md"Make sure that you return an array of objects of the type `Agent`.")
		elseif length(result) != 4
			almost(md"Make sure that you return `N` agents.")
		elseif length(Set(result)) != 4
			almost(md"You returned the **same** agent `N` times. You need to call the `Agent` constructor `N` times, not once.")
		else
			if sum(a -> a.status == I, result) != 1
				almost(md"Exactly one of the agents should be infectious.")
			else
				correct()
			end
		end
	end
end

# ╔═╡ 759bc42e-04ab-11eb-0ab1-b12e008c02a9
if !@isdefined(interact!)
	not_defined(:interact!)
else
	let
		agent = Agent(S, 9)
		source = Agent(I, 0)
		interact!(agent, source, InfectionRecovery(0.0, 1.0))

		if source.status != I || source.num_infected != 0
			keep_working(md"The `source` should not be modified if no infection occured.")
		elseif agent.status != S
			keep_working(md"The `agent` should get infected with the right probability.")
		else
			agent = Agent(S, 9)
			source = Agent(S, 0)
			interact!(agent, source, InfectionRecovery(1.0, 1.0))

			if source.status != S || source.num_infected != 0 || agent.status != S
				keep_working(md"The `agent` should get infected with the right probability if the source is infectious.")
			else
				agent = Agent(S, 9)
				source = Agent(I, 3)
				interact!(agent, source, InfectionRecovery(1.0, 1.0))

				if agent.status == R
					almost(md"The agent should not recover immediately after becoming infectious.")
				elseif agent.status == S
					keep_working(md"The `agent` should recover from an infectious state with the right probability.")
				elseif source.status != I || source.num_infected != 4
					almost(md"The `source` did not get updated correctly after infecting the `agent`.")
				else
					correct(md"Your function treats the **susceptible** agent case correctly!")
				end
			end
		end
	end
end

# ╔═╡ 1491a078-04aa-11eb-0106-19a3cf1e94b0
if !@isdefined(interact!)
	not_defined(:interact!)
else
	let
		agent = Agent(I, 9)
		source = Agent(S, 0)

		interact!(agent, source, InfectionRecovery(1.0, 1.0))

		if source.status != S || source.num_infected != 0
			keep_working(md"The `source` should not be modified if `agent` is infectious.")
		elseif agent.status != R
			keep_working(md"The `agent` should recover from an infectious state with the right probability.")
		elseif agent.num_infected != 9
			keep_working(md"`agent.num_infected` should not be modified if `agent` is infectious.")
		else
			let
				agent = Agent(I, 9)
				source = Agent(R, 0)

				interact!(agent, source, InfectionRecovery(1.0, 0.0))

				if agent.status == R
					keep_working(md"The `agent` should recover from an infectious state with the right probability.")
				else
					correct(md"Your function treats the **infectious** agent case correctly!")
				end
			end
		end
	end
end

# ╔═╡ f8e05d94-04ac-11eb-26d4-6f1d2c5ed272
if !@isdefined(interact!)
	not_defined(:interact!)
else
	let
		agent = Agent(R, 9)
		source = Agent(I, 0)
		interact!(agent, source, InfectionRecovery(1.0, 1.0))

		if source.status != I || source.num_infected != 0
			keep_working(md"The `source` should not be modified if no infection occured.")
		elseif agent.status != R || agent.num_infected != 9
			keep_working(md"The `agent` should not be momdified if it is in a recoved state.")
		else
			correct(md"Your function treats the **recovered** agent case correctly!")
		end
	end
end

# ╔═╡ 39dffa3c-0414-11eb-0197-e72b299e9c63
bigbreak = html"

";

# ╔═╡ 2b26dc42-0403-11eb-205f-cd2c23d8cb03
bigbreak

# ╔═╡ 5689841e-0414-11eb-0492-63c77ddbd136
bigbreak

# ╔═╡ Cell order:
# ╟─01341648-0403-11eb-2212-db450c299f35
# ╟─03a85970-0403-11eb-334a-812b59c0905b
# ╟─06f30b2a-0403-11eb-0f05-8badebe1011d
# ╠═095cbf46-0403-11eb-0c37-35de9562cebc
# ╟─107e65a4-0403-11eb-0c14-37d8d828b469
# ╠═12cc2940-0403-11eb-19a7-bb570de58f6f
# ╠═15187690-0403-11eb-2dfd-fd924faa3513
# ╟─1d3356c4-0403-11eb-0f48-01b5eb14a585
# ╟─2b26dc42-0403-11eb-205f-cd2c23d8cb03
# ╟─df8547b4-0400-11eb-07c6-fb370b61c2b6
# ╠═02b0c2fc-0415-11eb-2b40-7bca8ea4eef9
# ╟─b817f466-04d4-11eb-0a26-c1c667f9f7f7
# ╟─76d117d4-0403-11eb-05d2-c5ea47d06f43
# ╠═d57c6a5a-041b-11eb-3ab4-774a2d45a891
# ╠═4ad79def-f52b-4b6b-88fa-581a976fc605
# ╟─c61f35ea-04d6-11eb-2503-17a79f8d0298
# ╟─6d906d0c-0415-11eb-0c1c-b5c0aca841db
# ╟─6db6c894-0415-11eb-305a-c75b119d89e9
# ╟─6de37d6c-0415-11eb-1b05-85ac820016c7
# ╠═73047bba-0416-11eb-1047-23e9c3dbde05
# ╟─76f62d64-0403-11eb-27e2-3de58366b619
# ╠═c5c7cb86-041b-11eb-3360-45463105f3c9
# ╠═d8abd2f6-0416-11eb-1c2a-f9157d9760a7
# ╟─771c8f0c-0403-11eb-097e-ab24d0714ad5
# ╠═105d347e-041c-11eb-2fc8-1d9e5eda2be0
# ╠═1ca7a8c2-041a-11eb-146a-15b8cdeaea72
# ╟─08e2bc64-0417-11eb-1457-21c0d18e8c51
# ╟─77428072-0403-11eb-0068-81e3728f2ebe
# ╠═4b3ec86c-0419-11eb-26fd-cbbfdf19afa8
# ╟─dc784864-0430-11eb-1478-d1153e017310
# ╠═8a28c56e-04b4-11eb-279c-3b4dfb2a9f9b
# ╟─9374e63c-0493-11eb-0952-4b97512d7cdb
# ╠═823364ce-041c-11eb-2467-7ffa4f751527
# ╠═1ddbaa18-0494-11eb-1fc8-250ab6ae89f1
# ╟─f3f81172-041c-11eb-2b9b-e99b7b9400ed
# ╟─7768a2dc-0403-11eb-39b7-fd660dc952fe
# ╠═f1f89502-0494-11eb-2303-0b79d8bbd13f
# ╠═06089d1e-0495-11eb-0ace-a7a7dc60e5b2
# ╟─77b54c10-0403-11eb-16ad-65374d29a817
# ╠═103ed025-823f-4923-8e5f-55b2b694fe70
# ╠═f33f2b39-ed8a-4808-9422-1bd94718e76c
# ╠═bb63f3cc-042f-11eb-04ff-a128aec3c378
# ╟─bb8aeb58-042f-11eb-18b8-f995631df619
# ╠═992aadc6-4e4e-4f9c-9091-557d2a54bd25
# ╟─778ec25c-0403-11eb-3146-1d11c294bb1f
# ╟─5c9c0050-9f2a-41b6-9c71-b6fa4658330c
# ╠═f1693d9c-a6a7-4e0d-b958-1a6a459882d7
# ╠═b8a2f491-fc02-4ad1-a762-7d36643ce844
# ╠═7bb8e426-0495-11eb-3a8b-cbbab61a1631
# ╟─77db111e-0403-11eb-2dea-4b42ceed65d6
# ╠═321ac3d6-b40a-49d0-bbb5-5f1a9f6cd6dc
# ╠═3ea098eb-3dd2-4e1e-be87-8aad8ecc5dfb
# ╟─61789646-0403-11eb-0042-f3b8308f11ba
# ╠═26f84600-041d-11eb-1856-b12a3e5c1dc7
# ╟─271ec5f0-041d-11eb-041b-db46ec1465e0
# ╠═7f4e121c-041d-11eb-0dff-cd0cbfdfd606
# ╟─7f744644-041d-11eb-08a0-3719cc0adeb7
# ╠═88c53208-041d-11eb-3b1e-31b57ba99f05
# ╟─847d0fc2-041d-11eb-2864-79066e223b45
# ╠═cc8a6234-27e0-49c7-af1a-881d98a85ec6
# ╟─860790fc-0403-11eb-2f2e-355f77dcc7af
# ╠═ae4ac4b4-041f-11eb-14f5-1bcde35d18f2
# ╟─ae70625a-041f-11eb-3082-0753419d6d57
# ╠═60a8b708-04c8-11eb-37b1-3daec644ac90
# ╟─189cae1e-0424-11eb-2666-65bf297d8bdd
# ╠═18d308c4-0424-11eb-176d-49feec6889cf
# ╟─190deebc-0424-11eb-19fe-615997093e14
# ╟─b329070b-f735-47c7-84fa-d0f60dbdd666
# ╠═82f2580a-04c8-11eb-1eea-bdb4e50eee3b
# ╟─8631a536-0403-11eb-0379-bb2e56927727
# ╠═98beb336-0425-11eb-3886-4f8cfd210288
# ╠═26d27f54-bb4f-45e2-9e67-6f54621c3f18
# ╟─7c515a7a-04d5-11eb-0f36-4fcebff709d5
# ╟─866299e8-0403-11eb-085d-2b93459cc141
# ╠═9a837b52-0425-11eb-231f-a74405ff6e23
# ╠═a8dd5cae-0425-11eb-119c-bfcbf832d695
# ╟─c4a8694a-04d4-11eb-1eef-c9e037e6b21f
# ╟─8692bf42-0403-11eb-191f-b7d08895274f
# ╠═7946d83a-04a0-11eb-224b-2b315e87bc84
# ╠═488771e2-049f-11eb-3b0a-0de260457731
# ╟─393041ec-049f-11eb-3089-2faf378445f3
# ╟─86d98d0a-0403-11eb-215b-c58ad721a90b
# ╠═223933a4-042c-11eb-10d3-852229f25a35
# ╠═1a654bdc-0421-11eb-2c38-7d35060e2565
# ╟─2d3bba2a-04a8-11eb-2c40-87794b6aeeac
# ╠═406aabea-04a5-11eb-06b8-312879457c42
# ╟─b21475c6-04ac-11eb-1366-f3b5e967402d
# ╠═9c39974c-04a5-11eb-184d-317eb542452c
# ╟─759bc42e-04ab-11eb-0ab1-b12e008c02a9
# ╟─1491a078-04aa-11eb-0106-19a3cf1e94b0
# ╟─f8e05d94-04ac-11eb-26d4-6f1d2c5ed272
# ╟─619c8a10-0403-11eb-2e89-8b0974fb01d0
# ╠═2ade2694-0425-11eb-2fb2-390da43d9695
# ╟─955321de-0403-11eb-04ce-fb1670dfbb9e
# ╠═46133a74-04b1-11eb-0b46-0bc74e564680
# ╟─95771ce2-0403-11eb-3056-f1dc3a8b7ec3
# ╠═887d27fc-04bc-11eb-0ab9-eb95ef9607f8
# ╠═b92f1cec-04ae-11eb-0072-3535d1118494
# ╠═2c62b4ae-04b3-11eb-0080-a1035a7e31a2
# ╠═c5156c72-04af-11eb-1106-b13969b036ca
# ╟─28db9d98-04ca-11eb-3606-9fb89fa62f36
# ╟─0a967f38-0493-11eb-0624-77e40b24d757
# ╟─bf6fd176-04cc-11eb-008a-2fb6ff70a9cb
# ╠═38b1aa5a-04cf-11eb-11a2-930741fc9076
# ╠═80c2cd88-04b1-11eb-326e-0120a39405ea
# ╟─80e6f1e0-04b1-11eb-0d4e-475f1d80c2bb
# ╠═ca8e82da-6d97-4f02-9c66-8fd63349458b
# ╟─9cf9080a-04b1-11eb-12a0-17013f2d37f5
# ╟─95c598d4-0403-11eb-2328-0175ed564915
# ╠═843fd63c-04d0-11eb-0113-c58d346179d6
# ╠═7f635722-04d0-11eb-3209-4b603c9e843c
# ╟─dfb99ace-04cf-11eb-0739-7d694c837d59
# ╠═1c6aa208-04d1-11eb-0b87-cf429e6ff6d0
# ╠═19f75d2b-e35c-4475-a7b2-bc9d224075dc
# ╟─95eb9f88-0403-11eb-155b-7b2d3a07cff0
# ╠═287ee7aa-0435-11eb-0ca3-951dbbe69404
# ╠═2efb95c6-0f45-4c6a-a808-a4046daea0f8
# ╟─9611ca24-0403-11eb-3582-b7e3bb243e62
# ╠═26e2978e-0435-11eb-0d61-25f552d2771e
# ╠═c72e0b01-8730-4b24-aa1a-ba6e54326ad5
# ╟─f9cff122-d647-4a3f-861a-4a5967922368
# ╟─9635c944-0403-11eb-3982-4df509f6a556
# ╠═4ad11052-042c-11eb-3643-8b2b3e1269bc
# ╟─61c00724-0403-11eb-228d-17c11670e5d1
# ╠═8dd97820-04a5-11eb-36c0-8f92d4b859a8
# ╟─99ef7b2a-0403-11eb-08ef-e1023cd151ae
# ╟─9a13b17c-0403-11eb-024f-9b37e95e211b
# ╠═1ac4b33a-0435-11eb-36f8-8f3f81ae7844
# ╟─9a377b32-0403-11eb-2799-e7e59caa6a45
# ╠═21c50840-0435-11eb-1307-7138ecde0691
# ╟─da49710e-0420-11eb-092e-4f1173868738
# ╠═e6219c7c-0420-11eb-3faa-13126f7c8007
# ╟─5689841e-0414-11eb-0492-63c77ddbd136
# ╟─531d13c2-0414-11eb-0acd-4905a684869d
# ╟─4f19e872-0414-11eb-0dfd-e53d2aecc4dc
# ╟─48a16c42-0414-11eb-0e0c-bf52bbb0f618
# ╟─461586dc-0414-11eb-00f3-4984b57bfac5
# ╟─43e6e856-0414-11eb-19ca-07358aa8b667
# ╟─41cefa68-0414-11eb-3bad-6530360d6f68
# ╟─3f5e0af8-0414-11eb-34a7-a71e7aaf6443
# ╟─3d88c056-0414-11eb-0025-05d3aff1588b
# ╟─3c0528a0-0414-11eb-2f68-a5657ab9e73d
# ╟─39dffa3c-0414-11eb-0197-e72b299e9c63


text:
# This file is a part of Julia. License is MIT: http://julialang.org/license

# Convert # of Rata Die days to proleptic Gregorian calendar y,m,d,w
# Reference: http://mysite.verizon.net/aesir_research/date/date0.htm
function yearmonthday(days)
    z = days + 306; h = 100z - 25; a = fld(h,3652425); b = a - fld(a,4);
    y = fld(100b+h,36525); c = b + z - 365y - fld(y,4); m = div(5c+456,153);
    d = c - div(153m-457,5); return m > 12 ? (y+1,m-12,d) : (y,m,d)
end
function year(days)
   z = days + 306; h = 100z - 25; a = fld(h,3652425); b = a - fld(a,4);
   y = fld(100b+h,36525); c = b + z - 365y - fld(y,4); m = div(5c+456,153);
   return m > 12 ? y+1 : y
end
function yearmonth(days)
    z = days + 306; h = 100z - 25; a = fld(h,3652425); b = a - fld(a,4);
    y = fld(100b+h,36525); c = b + z - 365y - fld(y,4); m = div(5c+456,153);
    return m > 12 ? (y+1,m-12) : (y,m)
end
function month(days)
    z = days + 306; h = 100z - 25; a = fld(h,3652425); b = a - fld(a,4);
    y = fld(100b+h,36525); c = b + z - 365y - fld(y,4); m = div(5c+456,153);
    return m > 12 ? m-12 : m
end
function monthday(days)
    z = days + 306; h = 100z - 25; a = fld(h,3652425); b = a - fld(a,4);
    y = fld(100b+h,36525); c = b + z - 365y - fld(y,4); m = div(5c+456,153);
    d = c - div(153m-457,5); return m > 12 ? (m-12,d) : (m,d)
end
function day(days)
    z = days + 306; h = 100z - 25; a = fld(h,3652425); b = a - fld(a,4);
    y = fld(100b+h,36525); c = b + z - 365y - fld(y,4); m = div(5c+456,153);
    return c - div(153m-457,5)
end
# https://en.wikipedia.org/wiki/Talk:ISO_week_date#Algorithms
function week(days)
    w = div(abs(days-1),7) % 20871
    c,w = divrem((w + (w >= 10435)),5218)
    w = (w*28+[15,23,3,11][c+1]) % 1461
    return div(w,28) + 1
end

# Accessor functions
value(dt::TimeType) = dt.instant.periods.value
days(dt::Date) = value(dt)
days(dt::DateTime) = fld(value(dt),86400000)
year(dt::TimeType) = year(days(dt))
month(dt::TimeType) = month(days(dt))
week(dt::TimeType) = week(days(dt))
day(dt::TimeType) = day(days(dt))
hour(dt::DateTime)   = mod(fld(value(dt),3600000),24)
minute(dt::DateTime) = mod(fld(value(dt),60000),60)
second(dt::DateTime) = mod(fld(value(dt),1000),60)
millisecond(dt::DateTime) = mod(value(dt),1000)

dayofmonth(dt::TimeType) = day(dt)
yearmonth(dt::TimeType) = yearmonth(days(dt))
monthday(dt::TimeType) = monthday(days(dt))
yearmonthday(dt::TimeType) = yearmonthday(days(dt))

@vectorize_1arg TimeType year
@vectorize_1arg TimeType month
@vectorize_1arg TimeType day
@vectorize_1arg TimeType week
@vectorize_1arg DateTime hour
@vectorize_1arg DateTime minute
@vectorize_1arg DateTime second
@vectorize_1arg DateTime millisecond

@vectorize_1arg TimeType dayofmonth
@vectorize_1arg TimeType yearmonth
@vectorize_1arg TimeType monthday
@vectorize_1arg TimeType yearmonthday

max_stars_repo_name: lidiapopescu/JuliaDT

max_issues_repo_name: lidiapopescu/JuliaDT

max_forks_repo_name: JuliaComputing/JuliaDT


text:
# Routines for interaction with pymatgen, e.g. converting to
# its structures from the DFTK equivalents

function pymatgen_lattice(lattice::AbstractArray)
    # Notice: Pymatgen uses rows as lattice vectors, so we unpeel
    # our lattice column by column. The default unit in pymatgen is Ǎngström
    mg = pyimport("pymatgen")
    bohr_to_A = 1 / units.Ǎ
    mg.Lattice([Array(bohr_to_A .* lattice[:, i]) for i in 1:3])
end
pymatgen_lattice(model::Model) = pymatgen_lattice(model.lattice)

function pymatgen_structure(model_or_lattice, atoms)
    mg = pyimport("pymatgen")
    pylattice = pymatgen_lattice(model_or_lattice)

    n_species = sum(length(pos) for (spec, pos) in atoms)
    pyspecies = Vector{Int}(undef, n_species)
    pypositions = Array{Vector{Float64}}(undef, n_species)
    ispec = 1
    for (spec, pos) in atoms
        for coord in pos
            pyspecies[ispec] = spec.Z
            pypositions[ispec] = Vector{Float64}(coord)
            ispec = ispec + 1
        end
    end
    @assert ispec == n_species + 1

    mg.Structure(pylattice, pyspecies, pypositions)
end
pymatgen_structure(model::Model) = pymatgen_structure(model, model.atoms)

function load_lattice_pymatgen(T, pyobj::PyObject)
    mg = pyimport("pymatgen")

    if pyisinstance(pyobj, mg.Structure)
        load_lattice_pymatgen(T, pyobj.lattice)
    elseif pyisinstance(pyobj, mg.Lattice)
        lattice = Matrix{T}(undef, 3, 3)
        for i in 1:3, j in 1:3
            lattice[i, j] = units.Ǎ * get(get(pyobj.matrix, j-1), i-1)
        end
        Mat3{T}(lattice)
    else
        error("load_lattice_pymatgen not implemented for python type $pyobj")
    end
end

"""
Load a DFTK-compatible atoms representation from a supported pymatgen object.
All atoms are using a Coulomb model.
"""
function load_atoms_pymatgen(T, pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("pymatgen").Structure)
    map(unique(pyobj.species)) do spec
        coords = [s.frac_coords for s in pyobj.sites if s.specie == spec]
        ElementCoulomb(spec.number) => coords
    end
end


text:
#
# @defcomp and supporting functions
#
using MacroTools

# Store a list of built-in components so we can suppress messages about creating them.
# TBD: suppress returning these in the list of components at the user level.
const global built_in_comps = (:adder,  :multiplier, :ConnectorCompVector, :ConnectorCompMatrix)

is_builtin(comp_name) = comp_name in built_in_comps

function _generate_run_func(comp_name, module_name, args, body)
    if length(args) != 4
        error("Cannot generate run_timestep; requires 4 arguments but got $args")
    end

    (p, v, d, t) = args

    # Generate unique function name for each component so we can store a function pointer.
    # (Methods requiring dispatch cannot be invoked directly. Could use FunctionWrapper here...)
    func_name = Symbol("run_timestep_$(module_name)_$(comp_name)")

    # Needs "global" so function is defined outside the "let" statement
    func = :(
        global function $(func_name)($(p), #::Mimi.ComponentInstanceParameters,
                                     $(v), #::Mimi.ComponentInstanceVariables
                                     $(d), #::NamedTuple
                                     $(t)) #::T <: Mimi.AbstractTimestep
            $(body...)
            return nothing
        end
    )
    return func
end

function _generate_init_func(comp_name, module_name, args, body)

    if length(args) != 3
        error("Cannot generate init function; requires 3 arguments but got $args")
    end

    # add types to the parameters
    (p, v, d) = args

    func_name = Symbol("init_$(module_name)_$(comp_name)")

    func = :(
        global function $(func_name)($(p), #::Mimi.ComponentInstanceParameters
                                     $(v), #::Mimi.ComponentInstanceVariables
                                     $(d)) #::NamedTuple
            $(body...)
            return nothing
        end
    )
    @debug "init func: $func"
    return func
end

function _check_for_known_argname(name)
    if !(name in (:description, :unit, :index, :default))
        error("Unknown argument name: '$name'")
    end
end

function _check_for_known_element(name)
    if !(name in (:Variable, :Parameter, :Index))
        error("Unknown component element type: '$name'")
    end
end

# Add a variable to a ComponentDef. CompositeComponents have no vars of their own,
# only references to vars in components contained within.
function add_variable(comp_def::ComponentDef, name, datatype, dimensions, description, unit)
    v = VariableDef(name, comp_def.comp_path, datatype, dimensions, description, unit)
    comp_def[name] = v            # adds to namespace and checks for duplicate
    return v
end

# Add a variable to a ComponentDef referenced by ComponentId
function add_variable(comp_id::ComponentId, name, datatype, dimensions, description, unit)
    add_variable(compdef(comp_id), name, datatype, dimensions, description, unit)
end

function add_parameter(comp_def::ComponentDef, name, datatype, dimensions, description, unit, default)
    if default !== nothing
        ndims_default = (default isa Symbol || default isa String) ? 0 : ndims(default)
        if length(dimensions) != ndims_default
            error("Default value has different number of dimensions ($(ndims_default))) than parameter '$name' ($(length(dimensions)))")
        end
    end
    p = ParameterDef(name, comp_def.comp_path, datatype, dimensions, description, unit, default)
    comp_def[name] = p            # adds to namespace and checks for duplicate
    dirty!(comp_def)
    return p
end

function add_parameter(comp_id::ComponentId, name, datatype, dimensions, description, unit, default)
    add_parameter(compdef(comp_id), name, datatype, dimensions, description, unit, default)
end

# Generates an expression to construct a Variable or Parameter
function _generate_var_or_param(elt_type, name, datatype, dimensions, dflt, desc, unit)
    func_name = elt_type == :Parameter ? :add_parameter : :add_variable
    args = [datatype, dimensions, desc, unit]
    if elt_type == :Parameter
        push!(args, dflt)
    end

    expr = :(Mimi.$func_name(comp, $(QuoteNode(name)), $(args...)))
    @debug "Returning: $expr\n"
    return expr
end

function _generate_dims_expr(name, args, datum_type)
    @debug "  Index $name"

    # Args are not permitted; we attempt capture only to check syntax
    if length(args) > 0
        error("Index $name: arguments to Index() are not permitted")
    end

    # Ditto types for Index, e.g., region::Foo = Index()
    if datum_type !== nothing
        error("Index $name: Type specification ($datum_type) is not supported")
    end

    name_expr = (name isa Symbol) ? :($(QuoteNode(name))) : name
    expr = :(Mimi.add_dimension!(comp, $name_expr))
    @debug "  dims expr: $expr"
    return expr
end

_generate_dims_expr(name) = _generate_dims_expr(name, [], nothing)

"""
    defcomp(comp_name::Symbol, ex::Expr)

Define a Mimi component `comp_name` with the expressions in `ex`.  The following
types of expressions are supported:

1. `dimension_name = Index()`   # defines a dimension
2. `parameter = Parameter(index = [dimension_name], units = "unit_name", default = default_value)`    # defines a parameter with optional arguments
3. `variable = Variable(index = [dimension_name], units = "unit_name")`    # defines a variable with optional arguments
4. `init(p, v, d)`              # defines an init function for the component
5. `run_timestep(p, v, d, t)`   # defines a run_timestep function for the component

Parses a @defcomp definition, converting it into a series of function calls that
create the corresponding ComponentDef instance. At model build time, the ModelDef
(including its ComponentDefs) will be converted to a runnable model.
"""
macro defcomp(comp_name, ex)
    known_dims = Set{Union{Int, Symbol}}()

    @capture(ex, elements__)
    @debug "Component $comp_name"

    # TBD: Allow explicit definition of module to define component in
    if @capture(comp_name, module_name_.cmpname_)       # e.g., Mimi.adder
        comp_name = cmpname
    end

    # We'll return a block of expressions that will define the component.
    # N.B. @__MODULE__ is evaluated when the expanded macro is interpreted.
    result = :(
        let calling_module = @__MODULE__,
            comp_id = Mimi.ComponentId(calling_module, $(QuoteNode(comp_name))),
            comp = Mimi.ComponentDef(comp_id)

            global $comp_name = comp
        end
    )

    # helper function used in loop below
    function addexpr(expr)
        let_block = result.args[end].args
        push!(let_block, expr)
    end

    for elt in elements
        @debug "elt: $elt"

        # handle doc strings, which appear as macro calls to
        if @capture(elt, @name_ doc_String expr_) && name isa GlobalRef && name.name == Symbol("@doc")
            @debug "ignoring doc string: $doc"
            elt = expr  # extract the expression; ignore the doc string
        end

        if @capture(elt, function fname_(args__) body__ end)
            if fname == :run_timestep
                body = elt.args[2].args  # replace captured body with this, which includes line numbers
                expr = _generate_run_func(comp_name, nameof(__module__), args, body)

            elseif fname == :init
                body = elt.args[2].args  # as above
                expr = _generate_init_func(comp_name, nameof(__module__), args, body)
            else
                error("@defcomp can contain only these functions: init(p, v, d) and run_timestep(p, v, d, t)")
            end

            addexpr(expr)
            continue
        end

        # DEPRECATION
        if @capture(elt, name_::datum_type_ = elt_type_(args__))
            error("The following syntax has been deprecated in @defcomp: \"$name::$datum_type = $elt_type(...)\". Use curly bracket syntax instead: \"$name = $elt_type{$datum_type}(...)\"")
        elseif ! @capture(elt, name_ = (elt_type_{datum_type_}(args__) | elt_type_(args__)))
        end

        # elt_type is one of {:Variable, :Parameter, :Index}
        if elt_type == :Index
            expr = _generate_dims_expr(name, args, datum_type)
            push!(known_dims, name)
            addexpr(expr)

        elseif elt_type in (:Variable, :Parameter)
            @debug "  $elt_type $name"
            desc = ""
            unit = ""
            dflt = nothing
            dimensions = Array{Symbol}(undef, 0)

            for arg in args
                @debug "    arg: $arg"
                if @capture(arg, argname_ = value_)
                    _check_for_known_argname(argname)

                elseif @capture(arg, argname_)
                    # If it's an unknown arg, report *that* rather than the missing value
                    _check_for_known_argname(argname)
                    error("Argument '$arg' of $elt_type $name is missing a value")

                else
                    error("Badly formatted argument: $arg")
                end

                if @capture(arg, description = value_)
                    desc = value

                elseif @capture(arg, unit = value_)
                    unit = value

                elseif @capture(arg, index = [dims__])
                    @debug "    dims: $dims"

                    if !isempty(filter(x -> !(x isa Union{Int,Symbol}), dims))
                        error("Dimensions ($dims) must be defined by a Symbol placeholder or an Int")
                    end

                    append!(dimensions, map(Symbol, dims))  # converts, e.g., 4 into Symbol("4")

                    # Add undeclared dimensions on-the-fly
                    for dim in dims
                        if ! (dim in known_dims)
                            addexpr(_generate_dims_expr(dim))
                            push!(known_dims, dim)
                        end
                    end

                elseif @capture(arg, default = dflt_)
                    if elt_type == :Variable
                        error("Default values are permitted only for Parameters, not for Variables")
                    end
                    @debug "Default for parameter $name is $dflt"
                end
            end

            @debug "    index $(Tuple(dimensions)), unit '$unit', desc '$desc'"

            datum_type = (datum_type === nothing ? :Number : datum_type)
            addexpr(_generate_var_or_param(elt_type, name, datum_type, dimensions, dflt, desc, unit))
        else
            error("Unrecognized element type: $elt_type")
        end
    end

    addexpr(:(nothing))         # reduces noise
    return esc(result)
end


text:
abstract AbstractPhenotype{G <: AbstractGenome}

genome(pt::AbstractPhenotype) = error()
fitness(pt::AbstractPhenotype) = error()
crossover{T<:AbstractPhenotype}(gt1::T, gt2::T, o) = T(gt1, crossover(genome(gt1), genome(gt2), o))

function mutate!{T<:AbstractPhenotype}(
        pop::AbstractVector{T},
        o::Elitism)
    sort!(pop, rev = true)
    n_top = floor(Int, o.fraction * length(pop))
    @inbounds for i = (n_top+1):length(pop)
        gt = genome(pop[i])
        mutate!(gt, o.kind)
    end
    pop
end

function crossover!{T<:AbstractPhenotype}(
        pop::AbstractVector{T},
        o::Elitism)
    sort!(pop, rev = true)
    n_top = floor(Int, o.fraction * length(pop))
    @inbounds for i = (n_top+1):length(pop)
        i1 = rand(1:n_top)
        i2 = rand(1:n_top)
        pop[i] = crossover(pop[i1], pop[i2], o.kind)
    end
    pop
end

function worst{T<:AbstractPhenotype}(pop::AbstractVector{T})
    worst_pt = pop[1]
    worst_val = fitness(worst_pt)
    @inbounds for i in 2:length(pop)
        cur_val = fitness(pop[i])
        if cur_val < worst_val
            worst_val = cur_val
            worst_pt = pop[i]
        end
    end
    worst_pt
end

function best{T<:AbstractPhenotype}(pop::AbstractVector{T})
    best_pt = pop[1]
    best_val = fitness(best_pt)
    @inbounds for i in 2:length(pop)
        cur_val = fitness(pop[i])
        if cur_val > best_val
            best_val = cur_val
            best_pt = pop[i]
        end
    end
    best_pt
end


text:
module Docs

include("Helper.jl")
include("Analyzer.jl")
include("Matching.jl")
include("Destructuring.jl")
include("PatternStructure.jl")

end


text:
const DIRECTIVE = 1
const WHITESPACE = 2
const TEXT = 3


text:
using DelimitedFiles, SharedArrays, Random, LinearAlgebra, Clustering, Base
using DataFrames, CSV
using Statistics, ArgParse

nmats = 100

function max_eigv(gene) #returns largest eigenvalue of a randomized matrix generated from data.
    gran = zeros(size(gene))
    for i = 1:size(gene,2)
        gran[:,i] = gene[shuffle(1:end),i]
    end
 #    @time mr = corr_mat(gran)
    # println(mr[[end-1],end])
    mr = cor(transpose(gran))
    # println(mr[[end-1],end])
    return eigvals(mr)
end

parsed_args = Dict([("c", true), ("expfile", "Data/clean/Basal-clean.tsv")])

##============ 0.5 - Input
#! gname s1 s2 ... gstart chromname
df = CSV.read(parsed_args["expfile"], delim='\t')

df = df[df.chromname .== "17", :]
gene = convert(Array{Float64,2},df[:, 2:end-2])

evs = SharedArray{Float64, 2}((nmats,size(gene)[1]))
@time for i in 1:nmats
    @time evs[i,:] = max_eigv(gene) # return array of ngenes size
end
# All randomized eigenvalues
new_evs = vec(evs)
new_evs = filter(x->x>0.0000001, new_evs) #Drop Zeros

# Original eigenvalues
mr = cor(transpose(gene))
mr_e = eigvals(mr)
mr_e = filter(x->x>0.0000001, mr_e) #Drop Zeros

using PyPlot

#Figura 2a
plt.clf()
# color = '0.75'
# color = '#eeefff'
plt.hist(new_evs,normed=true,bins=100,label="Randomized",color = [:C1])
# grid("on")
xlabel("eigenvalue")
ylabel("Probability density")
legend()
ax = gca()
ax.spines["top"].set_visible(false)
ax.spines["right"].set_visible(false)
plot!(size=(1200,800))
savefig("Basal-ch17-eigendist.png")

#Figura 2b
plt.clf()
plt.hist(mr_e,normed=true,bins=100,label="Empirical data")
plt.hist(new_evs,normed=true,bins=100,alpha=0.7,label="Randomized")
# grid("on")
xlabel("eigenvalue")
ylabel("Probability density")
legend()
ax = gca()
ax.spines["top"].set_visible(false)
ax.spines["right"].set_visible(false)

# using GR
# histogram(evs)
# readline()

mr = cor(transpose(gene))
mr_e = eigvals(mr)

max_stars_repo_name: josemaz/gene-matrices

max_issues_repo_name: josemaz/gene-matrices

max_forks_repo_name: josemaz/gene-matrices


text:
# ---
# title: 228. Summary Ranges
# id: problem228
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/summary-ranges/description/>
# hidden: true
# ---
#
# You are given a **sorted unique** integer array `nums`.
#
# Return _the **smallest sorted** list of ranges that **cover all the numbers in
# the array exactly**_. That is, each element of `nums` is covered by exactly
# one of the ranges, and there is no integer `x` such that `x` is in one of the
# ranges but not in `nums`.
#
# Each range `[a,b]` in the list should be output as:
#
#   * `"a->b"` if `a != b`
#   * `"a"` if `a == b`
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [0,1,2,4,5,7]
#     Output: ["0->2","4->5","7"]
#     Explanation: The ranges are:
#     [0,2] --> "0->2"
#     [4,5] --> "4->5"
#     [7,7] --> "7"
#
#
# **Example 2:**
#
#
#
#     Input: nums = [0,2,3,4,6,8,9]
#     Output: ["0","2->4","6","8->9"]
#     Explanation: The ranges are:
#     [0,0] --> "0"
#     [2,4] --> "2->4"
#     [6,6] --> "6"
#     [8,9] --> "8->9"
#
#
# **Example 3:**
#
#
#
#     Input: nums = []
#     Output: []
#
#
# **Example 4:**
#
#
#
#     Input: nums = [-1]
#     Output: ["-1"]
#
#
# **Example 5:**
#
#
#
#     Input: nums = [0]
#     Output: ["0"]
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 20`
#   * `-231 <= nums[i] <= 231 - 1`
#   * All the values of `nums` are **unique**.
#   * `nums` is sorted in ascending order.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end