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

using Markdown
using InteractiveUtils

# ╔═╡ f11023e5-8f7b-4f40-86d3-3407b61863d9
using PlutoUI, Viznet, Compose, Plots

# ╔═╡ ce44f8bd-692e-4eab-9ba4-055b25e40c81
using ForwardDiff: Dual

# ╔═╡ a1ef579e-4b66-4042-944e-7e27c660095e
md"""
```math
\newcommand{\comment}[1]{{\bf  \color{blue}{\text{◂~ #1}}}}
```
"""

# ╔═╡ 9d11e058-a7d0-11eb-1d78-6592ff7a1b43
md"# An introduction to automatic differentiation

-- GiggleLiu"

# ╔═╡ b73157bf-1a77-47b8-8a06-8d6ec2045023
html"<button onclick='present()'>present</button>"

# ╔═╡ ec13e0a9-64ff-4f66-a5a6-5fef53428fa1
md"""
* What is automatic differentiation (AD)?
    * A true history of AD
    * Forward mode AD
    * Reverse mode AD
        * primitves on tensors (including tensorflow, pytorch et al.)
        * primitves on elementary instructions (usually source code transformation based)
        * defined on a reversible program
* Some applications in **scientific computing**
    * solving the graph embedding problem
    * inverse engineering a hamiltonian
    * obtaining maximum independent set (MIS) configurations
    * towards differentiating `expmv` ``\comment{will be used in our emulator}``
"""

# ╔═╡ f8b0d1ce-99f7-4729-b46e-126da540cbbe
md"""
## The true history of automatic differentiation
"""

# ╔═╡ 435ac19e-1c0c-4ee5-942d-f2a97c8c4d80
md"""
* 1964 ~ Robert Edwin Wengert, A simple automatic derivative evaluation program. ``\comment{first forward mode AD}``
* 1970 ~ Seppo Linnainmaa, Taylor expansion of the accumulated rounding error. ``\comment{first backward mode AD}``
* 1986 ~ Rumelhart, D. E., Hinton, G. E., and Williams, R. J., Learning representations by back-propagating errors.
* 1992 ~ Andreas Griewank, Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. ``\comment{foundation of source code transformation based AD.}``
* 2000s ~ The boom of tensor based AD frameworks for machine learning.
* 2018 ~ People re-invented AD as differential programming ([wiki](https://en.wikipedia.org/wiki/Differentiable_programming) and this [quora answer](https://www.quora.com/What-is-Differentiable-Programming).)
![](https://qph.fs.quoracdn.net/main-qimg-fb2f8470f2120eb49c8142b08d9c4132)
* 2020 ~ Me, Differentiate everything with a reversible embeded domain-specific language ``\comment{AD based on reversible programming}``.
"""

# ╔═╡ 48ecd619-d01d-43ff-8b52-7c2566c3fa2b
md"## Forward mode automatic differentiation"

# ╔═╡ 4878ce45-40ff-4fae-98e7-1be41e930e4d
md"""
Forward mode AD attaches a infitesimal number $\epsilon$ to a variable, when applying a function $f$, it does the following transformation
```math
\begin{align}
    f(x+g \epsilon) = f(x) + f'(x) g\epsilon + \mathcal{O}(\epsilon^2)
\end{align}
```

The higher order infinitesimal is ignored.

**In the program**, we can define a *dual number* with two fields, just like a complex number
```
f((x, g)) = (f(x), f'(x)*g)
```
"""

# ╔═╡ b2c1936c-2c27-4fbb-8183-e38c5e858483
res = sin(Dual(π/4, 2.0))

# ╔═╡ 8be1b812-fcac-404f-98aa-0571cb990f34
res === Dual(sin(π/4), cos(π/4)*2.0)

# ╔═╡ 33e0c762-c75e-44aa-bfe2-bff92dd1ace8
md"
We can apply this transformation consecutively, it reflects the chain rule.
```math
\begin{align}
\frac{\partial \vec y_{i+1}}{\partial x} &= \boxed{\frac{\partial \vec y_{i+1}}{\partial \vec y_i}}\frac{\partial \vec y_i}{\partial x}\\
&\text{local Jacobian}
\end{align}
```
"

# ╔═╡ c59c35ee-1907-4736-9893-e22c052150ca
let
	lb = textstyle(:math, fontsize(8), width=0.5, height=0.5)
	tb = textstyle(:default, fontsize(10), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(3.5), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("white"), Compose.stroke("black"); r=0.08)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))
	ebr = bondstyle(:default, Compose.stroke("red"), linewidth(0.5mm))
	ebd = bondstyle(:default, linewidth(0.5mm), dashed=true)
	eba = bondstyle(:default, linewidth(0.5mm), Compose.arrow(), Compose.stroke("red"), Compose.fill("red"))

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, fill("red"), θ=π/2-atan((y .- x)...)-1π/6)
		ebr >> (x, y)
		t >> mid
	end

	Compose.set_default_graphic_size(15cm, 5cm)
	x = (0.1, 0.5)
	fi0 = (0.35, 0.5)
	fi1 = (0.7, 0.5)
	fi2 = (1.0, 0.5)
	img = canvas() do
		nb >> fi0
		nb >> fi1
		lb >> (fi0 .- (0.05, 0.1), "f_{i-1}")
		lb >> (fi1 .- (0.02, 0.1), "f_{i}")
		lb >> (x, "x")
		lb >> ((fi1 .+ fi0) ./ 2 .- (0.02, 0.0), raw"\vec{y}_{i}")
		lb >> ((fi1 .+ fi2) ./ 2 .- (0.05, 0.0), raw"\vec{y}_{i+1}")
		lb >> ((fi1 .+ fi2) ./ 2 .- (0.05, 0.0), "\\vec{y}_{i+1}")
		lb >> (x .- (0.00, 0.25), raw"\color{red}{1}")
		lb >> ((fi1 .+ fi0) ./ 2 .- (0.05, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i}}{\partial x}}")
		lb >> ((fi1 .+ fi2) ./ 2 .- (0.08, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i+1}}{\partial x}}")
		ebd >> (x, fi0)
		eb >> (fi0, fi1)
		eb >> (fi1, fi2)
		#arrow((fi1 .+ fi0) ./ 2 .+ (0.08, -0.3), (fi1 .+ fi2) ./ 2 .+ (-0.08, -0.3))
		arrow((fi1 .+ fi0) ./ 2 .+ (0.08, -0.3), (fi1 .+ fi2) ./ 2 .+ (-0.08, -0.3))
	end
	img
end

# ╔═╡ f12b25d8-7c78-4686-b46d-00b34e565605
let
	x = Dual(π/4, 1.0)
	for i=1:10
		x = sin(x)
	end
	x
end

# ╔═╡ d90c3cc9-084d-4cf7-9db7-42cea043030b
md"""
**Example:** Computing two gradients $\frac{\partial z\sin x}{\partial x}$ and $\frac{\partial \sin^2x}{\partial x}$ at one sweep
"""

# ╔═╡ 93c98cb2-18af-47df-afb3-8c5a34b4723c
let
	lb = textstyle(:math, fontsize(8), width=1.0, height=0.5)
	tb = textstyle(:default, fontsize(3.5), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(4.5), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("black"), Compose.stroke("transparent"); r=0.05)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))

	x_x = (0.1, 0.25)
	x_y = (0.9, 0.5)
	x_y2 = (0.9, 0.25)
	x_z = (0.3, 0.5)
	x_sin = (0.3, 0.25)
	x_mul = (0.5, 0.5)
	x_square = (0.5, 0.25)

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, θ=π/2-atan((y .- x)...)-1π/6)
		eb >> (x, y)
		t >> mid
	end

	img = canvas() do
		nb >> x_sin
		nb >> x_mul
		nb >> x_square
		tb_big >> (x_sin, "sin")
		tb_big >> (x_mul .+ (0, 0.01), "*")
		tb_big >> (x_square, "^2")
		arrow(x_sin, x_mul)
		arrow(x_x, x_sin)
		arrow(x_mul, x_y)
		arrow(x_square, x_y2)
		arrow(x_z, x_mul)
		arrow(x_sin, x_square)
		tb >> ((x_x .+ x_sin) ./ 2 .- (0.02, 0.04), "x+ϵˣ")
		tb >> ((x_sin .+ x_mul) ./ 2 .- (0.08, 0.04), "sin(x)+cos(x)*ϵˣ")
		tb >> ((x_y .+ x_mul) ./ 2 .- (-0.04, 0.055), "z*sin(x)\n+z*cos(x)*ϵˣ")
		tb >> ((x_y2 .+ x_square) ./ 2 .- (-0.04, 0.055), "sin(x)^2\n+2*sin(x)*cos(x)*ϵˣ")
		tb >> ((x_z .+ x_mul) ./ 2 .- (0.05, 0.02), "z")
	end

	Compose.set_default_graphic_size(100mm, 100mm/2)
	Compose.compose(context(0, -0.15, 1, 2), img)
end

# ╔═╡ 2dc74e15-e2ea-4961-b43f-0ada1a73d80a
md"so the gradients are $z\cos x$ and $2\sin x\cos x$"

# ╔═╡ 7ee75a15-eaea-462a-92b6-293813d2d4d7
md"""
**What if we want to compute gradients for multiple inputs?**

The computing time grows **linearly** as the number of variables that we want to differentiate. But does not grow significantly with the number of outputs.
"""

# ╔═╡ 02a25b73-7353-43b1-8738-e7ca472d0cc7
md"""
## Reverse mode automatic differentiation

"""

# ╔═╡ 2afb984f-624e-4381-903f-ccc1d8a66a17
md"On the other side, the back-propagation can differentiate **many inputs** with respect to a **single output** efficiently"

# ╔═╡ 7e5d5e69-90f2-4106-8edf-223c150a8168
md"""
```math
\begin{align}
    \frac{\partial \mathcal{L}}{\partial \vec y_i} = \frac{\partial \mathcal{L}}{\partial \vec y_{i+1}}&\boxed{\frac{\partial \vec y_{i+1}}{\partial \vec y_i}}\\
&\text{local jacobian?}
\end{align}
```
"""

# ╔═╡ 92d7a938-9463-4eee-8839-0b8c5f762c79
let
	lb = textstyle(:math, fontsize(8), width=0.5, height=0.5)
	tb = textstyle(:default, fontsize(10), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(3.5), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("white"), Compose.stroke("black"); r=0.08)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))
	ebr = bondstyle(:default, Compose.stroke("red"), linewidth(0.5mm))
	ebd = bondstyle(:default, linewidth(0.5mm), dashed=true)
	eba = bondstyle(:default, linewidth(0.5mm), Compose.arrow(), Compose.stroke("red"), Compose.fill("red"))

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, fill("red"), θ=π/2-atan((y .- x)...)-1π/6)
		ebr >> (x, y)
		t >> mid
	end

	Compose.set_default_graphic_size(15cm, 5cm)
	x = (0.1, 0.5)
	fi0 = (0.35, 0.5)
	fi1 = (0.7, 0.5)
	fi2 = (0.9, 0.5)
	img = canvas() do
		nb >> fi0
		nb >> fi1
		lb >> (fi0 .- (0.02, 0.1), "f_{i}")
		lb >> (fi1 .- (0.05, 0.1), "f_{i+1}")
		lb >> (fi2 .- (0.05, 0.0), raw"\mathcal{L}")
		lb >> ((fi0 .+ x) ./ 2 .- (0.05, 0.0), raw"\vec{y}_{i}")
		lb >> ((fi0 .+ fi1) ./ 2 .- (0.05, 0.0), raw"\vec{y}_{i+1}")
		lb >> ((fi0 .+ fi1) ./ 2 .- (0.05, 0.0), "\\vec{y}_{i+1}")
		lb >> (fi2 .- (0.05, 0.25), raw"\color{red}{1}")
		lb >> ((fi0 .+ x) ./ 2 .- (0.08, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i}}{\partial x}}")
		lb >> ((fi0 .+ fi1) ./ 2 .- (0.08, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i+1}}{\partial x}}")
		ebd >> (fi1, fi2)
		eb >> (fi0, fi1)
		eb >> (x, fi0)
		#arrow((fi1 .+ fi0) ./ 2 .+ (0.08, -0.3), (fi1 .+ fi2) ./ 2 .+ (-0.08, -0.3))
		arrow( (fi0 .+ fi1) ./ 2 .+ (-0.08, -0.3), (fi0 .+ x) ./ 2 .+ (0.05, -0.3),)
	end
	img
end

# ╔═╡ 4b1a0b59-ddc6-4b2d-b5f5-d92084c31e46
md"### How to visite local Jacobians in the reversed order? "

# ╔═╡ a7fc71a2-6d45-4162-8073-4ddb85ded2e8
md"
**Design Decision**

1. Compute forward pass and caching inetermediate results into a global stack $\Sigma$ （packages except NiLang），
2. reversible programming."

# ╔═╡ fb6c3a48-550a-4d2e-a00b-a1e40d86b535
md"""
**Example:** Computing the gradient $\frac{\partial z\sin x}{\partial x}$ and $\frac{\partial z\sin x}{\partial z}$ by back propagating cached local information.
"""

# ╔═╡ ab6fa4ac-29ed-4722-88ed-fa1caf2072f3
let
	lb = textstyle(:math, fontsize(10), width=1.0, height=0.5)
	tb = textstyle(:default, fontsize(3.5), Compose.font("monospace"))
	tbc = textstyle(:default, fontsize(3.5), fill("red"), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(4), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("black"), Compose.stroke("transparent"); r=0.05)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))

	x_x = (0.1, 0.2)
	x_y = (0.9, 0.5)
	x_z = (0.1, 0.7)
	x_sin = (0.3, 0.3)
	x_mul = (0.5, 0.5)

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, θ=π/2-atan((y .- x)...)-1π/6)
		eb >> (x, y)
		t >> mid
	end
	img1 = canvas() do
		nb >> x_sin
		nb >> x_mul
		tb_big >> (x_sin, "sin")
		tb_big >> (x_mul .+ (0, 0.01), "*")
		arrow(x_sin, x_mul)
		arrow(x_x, x_sin)
		arrow(x_mul, x_y)
		arrow(x_z, x_mul)
		tb >> ((x_x .+ x_sin) ./ 2 .- (0.0, 0.1), "x \n push(Σ,x)")
		tb >> ((x_sin .+ x_mul) ./ 2 .- (-0.15, 0.04), "s = sin(x) \n push(Σ,s)")
		tb >> ((x_y .+ x_mul) ./ 2 .- (-0.05, 0.04), "y = z*sin(x)")
		tb >> ((x_z .+ x_mul) ./ 2 .- (0.05, 0.07), "z\n push(Σ,z)")
	end
	img2 = canvas() do
		nb >> x_sin
		nb >> x_mul
		tb_big >> (x_sin, "sin")
		tb_big >> (x_mul .+ (0, 0.01), "*")
		arrow(x_mul, x_sin)
		arrow(x_sin, x_x)
		arrow(x_y, x_mul)
		arrow(x_mul, x_z)
		tb >> ((x_x .+ x_sin) ./ 2 .- (0.0, 0.1), "x = pop(Σ)\nx̄ = cos(x)*s̄")
		tb >> ((x_sin .+ x_mul) ./ 2 .- (-0.12, 0.04), "z = pop(Σ)\ns̄ = z*ȳ")
		tb >> ((x_y .+ x_mul) ./ 2 .- (-0.05, 0.06), "y\nȳ=1")
		tb >> ((x_z .+ x_mul) ./ 2 .- (0.05, 0.07), "s = pop(Σ)\nz̄ = s*ȳ")
	end

	Compose.set_default_graphic_size(150mm, 75mm/1.4)
	Compose.compose(context(),
	(context(0, -0.1, 0.5, 1.4), img1),
	(context(0.5, -0.1, 0.5, 1.4), img2)
	)
end

# ╔═╡ 8e72d934-e307-4505-ac82-c06734415df6
md"Here, we use $\overline y$ for $\frac{\partial \mathcal{L}}{\partial y}$, which is also called the adjoint."

# ╔═╡ e6ff86a9-9f54-474b-8111-a59a25eda506
md"### Primitives on different scales"

# ╔═╡ 9c1d9607-a634-4350-aacd-2d40984d647d
md"We call the leaf nodes defining AD rules \"**primitives**\""

# ╔═╡ 63db2fa2-50b2-4940-b8ee-0dc6e3966a57
md"
**Design Decision**

* A: If we define primitives on **arrays**, we need tons of manually defined backward rules. (Jax, Pytorch, Zygote.jl, ReverseDiff.jl et al.)
* B: If we define primitives on **scalar instructions**, we will have worse tensor performance. (Tapenade, Adept, NiLang et al.)

*Note*: Here, implementing AD on scalars means specifically the **optimal checkpointing** approach, rather than a package like Jax, Zygote and ReverseDiff that having scalar support.
"

# ╔═╡ 693167e7-e80c-401d-af89-55b5fae30848
let
	w, h = 0.22, 0.1
	lb = Compose.compose(context(), polygon([(-w, -h), (-w, h), (w, h), (w, -h)]), Compose.stroke("transparent"))
	lb2 = Compose.compose(context(), polygon([(-w, -h), (-w, h), (w, h), (w, -h)]), Compose.stroke("transparent"), fill("red"))
	tb = Compose.compose(context(), Compose.text(0.0, 0.0, ""), fontsize(3), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(3), fill("white"), Compose.font("monospace"))
	eb = bondstyle(:default, linewidth(0.5mm))
	ar = bondstyle(:default, linewidth(0.3mm), Compose.arrow())
	xprog = (0.25, 0.15)
	xtensors = (0.25, 0.5)
	t1 = (0.5, 0.15)
	t2 = (0.5, 0.5)
	t3 = (0.5, 0.85)
	xscalars2 = (0.25, 0.85)

	function box(loc, text; color="black")
		(color=="black" ? lb : lb2) >> loc
		tb_big >> (loc, text)
	end
	Compose.set_default_graphic_size(10cm, 5cm)
	canvas() do
		box(xprog, "Program")
		ar >> (xprog, xtensors .+ (0, -h-0.03))
		#ar >> (xprog, xscalars .+ (-w/2, -h-0.03))
		ar >> (xtensors, xscalars2 .+ (0, -h-0.05))
		box(xtensors, "Functions on arrays")
		#box(xscalars, "Functions on Scalars")
		box(xscalars2, "Finite instructions"; color="red")
		tb >> (t1, "Neural networks")
		tb >> (t2, "matrix multiplication")
		tb >> (t3, "+, -, *")
	end
end

# ╔═╡ 4cd70901-2142-4868-9a33-c46ca0d064ec
html"""

on tensors
on finite instructions

meaning
defining backward rules manully for functions on tensors
defining backward rules on a limited set of basic scalar operations, and generate gradient code using source code transformation

pros and cons

Good tensor performance

Mature machine learning ecosystem

Need to define backward rules manually

Reasonalbe scalar performance

hard to utilize GPU kernels (except NiLang.jl) and BLAS

packages
Jax
PyTorch
Tapenade

Adept

NiLang.jl

"""

# ╔═╡ 4ff09f7c-aeac-48bd-9d58-8446137c3acd
md"""
## The AD ecosystem in Julia

Please check JuliaDiff: [https://juliadiff.org/](https://juliadiff.org/)

A short list:
* Forward mode AD: ForwardDiff.jl
* Reverse mode AD (tensor): ReverseDiff.jl/Zygote.jl
* Reverse mode AD (scalar): NiLang.jl

Warnings
* The main authors of `Tracker`, `ReverseDiff` and `Zygote` are not maintaining them anymore.
"""
#=
|       |   Rules | Favors Tensor? | Type |
| ---- | ---- | --- | --- |
|  Zygote   |  C  |  ✓   |   R     |
|  ReverseDiff  |  D    | ✓    | R |
|  Nabla   |  D→C  |   ✓  |   R     |
|  Tracker  |  D    | ✓    | R |
|  Yota   |  C  |  ✓   |     R   |
|  NiLang   |  -  |  ×   |  R      |
|  Enzyme   |  -  |  ×   |  R      |
|  ForwardDiff   |  -  |  ×   |    F    |
|  Diffractor   |  ?  |  ?   |  ?      |

* R: reverse mode
* F: forward mode
* C: ChainRules
* D: DiffRules
"""
=#

# ╔═╡ ea44037b-9359-4fbd-990f-529d88d54351
md"# Quick summary
1. The history of AD is longer than many people have thought. People are most familar with *reverse mode AD with primitives implemented on tensors* that brings the boom of machine learning. There are also AD frameworks that can differentiate a general program directly, which does not require users defining AD rules manually.
2. **Forward mode AD** propagate gradients forward, it has a computational overhead propotional to the number of input parameters.
2. **Backward mode AD** propagate gradients backward, it has a computational overhead propotional to the number of output parameters.
    * primitives on **tensors** v.s. **scalars**
    * reverse the program tape by **caching/checkpointing** v.s. **reversible programming**
4. Julia has one of the most active AD community!

#### Forward v.s. Backward
when is forward mode AD more useful?

* It is often combined with backward mode AD for obtaining Hessians (forward over backward).
* Having <20 input parameters.

when is backward mode AD more useful?
* In most variational optimizations, especially when we are training a neural network with ~ 100M parameters.
"

# ╔═╡ aa1547f2-5edd-4b7e-b93e-bdfc4e4fc6d5
md"""# The true reverse mode "automatic" differentiation"""

# ╔═╡ 71f4b476-027d-4c8f-b561-1ee418bc9e61
md"""
![](https://raw.githubusercontent.com/GiggleLiu/WuLiXueBao/master/paper/bennett_treeverse_pebbles.svg)
"""

# ╔═╡ ae096ad2-3ae9-4440-a959-0d7d9a174f1d
md"## How to differentiate sparse matrix multiplication"

# ╔═╡ 11557d6b-3a1e-416d-874f-b8d217976f76
md"## How to differentiate QR"

# ╔═╡ 6e7bf956-67b1-430c-9cf5-1746282727b2
md"## Time and space"

# ╔═╡ Cell order:
# ╟─a1ef579e-4b66-4042-944e-7e27c660095e
# ╠═f11023e5-8f7b-4f40-86d3-3407b61863d9
# ╟─9d11e058-a7d0-11eb-1d78-6592ff7a1b43
# ╟─b73157bf-1a77-47b8-8a06-8d6ec2045023
# ╟─ec13e0a9-64ff-4f66-a5a6-5fef53428fa1
# ╟─f8b0d1ce-99f7-4729-b46e-126da540cbbe
# ╟─435ac19e-1c0c-4ee5-942d-f2a97c8c4d80
# ╟─48ecd619-d01d-43ff-8b52-7c2566c3fa2b
# ╟─4878ce45-40ff-4fae-98e7-1be41e930e4d
# ╠═ce44f8bd-692e-4eab-9ba4-055b25e40c81
# ╠═b2c1936c-2c27-4fbb-8183-e38c5e858483
# ╠═8be1b812-fcac-404f-98aa-0571cb990f34
# ╟─33e0c762-c75e-44aa-bfe2-bff92dd1ace8
# ╟─c59c35ee-1907-4736-9893-e22c052150ca
# ╠═f12b25d8-7c78-4686-b46d-00b34e565605
# ╟─d90c3cc9-084d-4cf7-9db7-42cea043030b
# ╟─93c98cb2-18af-47df-afb3-8c5a34b4723c
# ╟─2dc74e15-e2ea-4961-b43f-0ada1a73d80a
# ╟─7ee75a15-eaea-462a-92b6-293813d2d4d7
# ╟─02a25b73-7353-43b1-8738-e7ca472d0cc7
# ╟─2afb984f-624e-4381-903f-ccc1d8a66a17
# ╟─7e5d5e69-90f2-4106-8edf-223c150a8168
# ╟─92d7a938-9463-4eee-8839-0b8c5f762c79
# ╟─4b1a0b59-ddc6-4b2d-b5f5-d92084c31e46
# ╟─a7fc71a2-6d45-4162-8073-4ddb85ded2e8
# ╟─fb6c3a48-550a-4d2e-a00b-a1e40d86b535
# ╟─ab6fa4ac-29ed-4722-88ed-fa1caf2072f3
# ╟─8e72d934-e307-4505-ac82-c06734415df6
# ╟─e6ff86a9-9f54-474b-8111-a59a25eda506
# ╟─9c1d9607-a634-4350-aacd-2d40984d647d
# ╟─63db2fa2-50b2-4940-b8ee-0dc6e3966a57
# ╟─693167e7-e80c-401d-af89-55b5fae30848
# ╟─4cd70901-2142-4868-9a33-c46ca0d064ec
# ╟─4ff09f7c-aeac-48bd-9d58-8446137c3acd
# ╟─ea44037b-9359-4fbd-990f-529d88d54351
# ╟─aa1547f2-5edd-4b7e-b93e-bdfc4e4fc6d5
# ╟─ae096ad2-3ae9-4440-a959-0d7d9a174f1d
# ╟─11557d6b-3a1e-416d-874f-b8d217976f76
# ╟─6e7bf956-67b1-430c-9cf5-1746282727b2
# ╟─71f4b476-027d-4c8f-b561-1ee418bc9e61


text:
function all_taxable()::Incomes_Dict
    eis = union(Set( keys( Exempt_Income )), Expenses )
    all_t = Incomes_Dict()
    for i in instances(Incomes_Type)
        if ! (i ∈ eis )
         all_t[i]=1.0
        end
    end
    all_t
end

const SAVINGS_INCOME :: Incomes_Dict = Incomes_Dict(

)

const DIVIDENDS = Incomes_Dict(
    dividends => 1.0
)

function make_non_savings()::Incomes_Dict
    excl = UnionAll(Set(Keys(DIVIDENDS)), SET( keys(SAVINGS_INCOME)))
    nsi = all_taxable()
    for i in excl
        delete!( nsi, i )
    end
    nsi
end

const NON_SAVINGS :: Incomes_Dict = make_non_savings()

ext: jl


text:
function prior()

end

function variational_bayes()

end

max_stars_repo_name: SosUts/spt

max_issues_repo_name: SosUts/spt

max_forks_repo_name: SosUts/spt


text:
using repeat_devlp
using Documenter

DocMeta.setdocmeta!(repeat_devlp, :DocTestSetup, :(using repeat_devlp); recursive=true)

makedocs(;
    modules=[repeat_devlp],
    authors="Gudongyangg",
    repo="https://github.com/Gudongyangg/repeat_devlp.jl/blob/{commit}{path}#{line}",
    sitename="repeat_devlp.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://Gudongyangg.github.io/repeat_devlp.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/Gudongyangg/repeat_devlp.jl",
    devbranch="master",
)


text:
module JuliaLearn

# PrintSourceCode
export PrintSourceCode
include("PrintSourceCode.jl")

# CombineScreenShots
export CombineScreenShots
include("CombineScreenShots.jl")

# PlotTopography
export PlotTopography
include("PlotTopography.jl")

end # module


text:
"""
    head_to_obj(filename::String)

Returns the header of a snapshot as a `Header` object.
"""
function head_to_obj(filename)

    h = Header()

    f = open(filename)
    blocksize = read(f, Int32)

    if blocksize == 8
        swap = false
        snap_format = 2
    elseif blocksize == 256
        swap = false
        snap_format = 1
    else
        blocksize = bswap(blocksize)
        if blocksize == 8
            swap = true
            snap_format = 2
        elseif blocksize == 256
            swap = true
            snap_format = 1
        else
            error("incorrect file format encountered when reading header of $filename")
        end
    end

    #println("Reading snapshot format: ", snap_format)

    if snap_format == 2
        seek(f, 16)
        skip_line = read(f, Int32)
    end

    h.npart = read!(f, Array{Int32,1}(undef,6))
    h.massarr = read!(f, Array{Float64,1}(undef,6))
    h.time = read(f, Float64)
    h.z = read(f, Float64)
    h.flag_sfr = read(f, Int32)
    h.flag_feedback = read(f, Int32)
    h.nall = read!(f, Array{UInt32,1}(undef,6))
    h.flag_cooling = read(f, Int32)
    h.num_files = read(f, Int32)
    h.boxsize = read(f, Float64)
    h.omega_0 = read(f, Float64)
    h.omega_l = read(f, Float64)
    h.h0 = read(f, Float64)
    h.flag_stellarage = read(f, Int32)
    h.flag_metals = read(f, Int32)
    h.npartTotalHighWord = read!(f, Array{UInt32,1}(undef,6))
    h.flag_entropy_instead_u = read(f, Int32)
    h.flag_doubleprecision = read(f, Int32)
    h.flag_ic_info = read(f, Int32)
    h.lpt_scalingfactor = read(f, Float32)

    close(f)

    return h
end

"""
    head_to_dict(filename::String)

Returns the header of a snapshot as a dictionary.
"""
function head_to_dict(filename::String)

        header = Dict()

        f = open(filename)
        blocksize = read(f, Int32)

        if blocksize[1] == 8
            swap = 0
            snap_format = 2
        elseif blocksize[1] == 256
            swap = 0
            snap_format = 1
        else
            blocksize[1] = bswap(blocksize[1])
            if blocksize[1] == 8
                swap = 1
                snap_format = 2
            elseif blocksize[1] == 256
                swap = 1
                snap_format = 1
            else
                println("incorrect file format encountered when reading header of", filename)
            end
        end

        if snap_format == 2
            seek(f, 16)
            skip_line = read(f, Int32)
        end

        header["snap_format"] = snap_format
        header["PartTypes"] = ["PartType0", "PartType1", "PartType2",
                               "PartType3", "PartType4", "PartType5"]
        header["npart"] = read!(f, Array{Int32,1}(undef,6))
        header["massarr"] = read!(f, Array{Float64,1}(undef,6))
        header["time"] = read(f, Float64)
        header["redshift"] = read(f, Float64)
        header["flag_sfr"] = read(f, Int32)
        header["flag_feedback"] = read(f, Int32)
        header["nall"] = read!(f, Array{UInt32,1}(undef,6))
        header["flag_cooling"] = read(f, Int32)
        header["num_files"] = read(f, Int32)
        header["boxsize"] = read(f, Float64)
        header["omega_m"] = read(f, Float64)
        header["omega_l"] = read(f, Float64)
        header["hubble"] = read(f, Float64)
        header["flag_stellarage"] = read(f, Int32)
        header["flag_metals"] = read(f, Int32)
        header["npartTotalHighWord"] = read!(f, Array{UInt32,1}(undef,6))
        header["flag_entropy_instead_u"] = read(f, Int32)
        header["flag_doubleprecision"] = read(f, Int32)
        header["flag_ic_info"] = read(f, Int32)
        header["lpt_scalingfactor"] = read(f, Float32)

        close(f)

        return header

end

"""
    read_header(filename::String)

Reads the header of a snapshot and returns a Header object.

See also: [`head_to_obj`](@ref)
"""
function read_header(filename::String)
    return head_to_obj(filename)
end

max_stars_repo_name: LudwigBoess/GadJet

max_issues_repo_name: LudwigBoess/GadJet

max_forks_repo_name: LudwigBoess/GadJet


text:
using MathepiaInference
using Test

@testset "MathepiaInference.jl" begin
    # Write your tests here.
end


text:
""" testing the viscosity routine
"""

import Particula as P
using Unitful, Test

@testset "µ" begin
    @test P.µ() |> upreferred |> unit == u"kg/m/s"
    @test P.µ(T = 300) |> unit == u"kg/m/s"
    @test_throws ArgumentError P.µ(T = 300 * u"C")
end


text:
function extractdocs(path::String)
    io = open(path, "r")

    inchunk = false
    chunks = Array(Vector{Char}, 0)
    i = 0
    if !eof(io)
        c = read(io, Char)#; print(c)
    end

    while !eof(io)
        if !inchunk
            if c == '#' && !eof(io)
                c = read(io, Char)#; print(c)
                if c == '\'' && !eof(io)
                    c = read(io, Char)#; print(c)
                    if c == ' '
                        i += 1
                        inchunk = true
                        push!(chunks, Array(Char, 0))
                    end
                end
            end
            if !eof(io)
                c = read(io, Char)#; print(c)
            end
        else
            if c != '\n'
                push!(chunks[i], c)
                if !eof(io)
                    c = read(io, Char)#; print(c)
                end
            else
                push!(chunks[i], c)
                if !eof(io)
                    c = read(io, Char)#; print(c)
                end
                while isspace(c) && c != '\n' && !eof(io)
                    c = read(io, Char)#; print(c)
                end
                if c == '#' && !eof(io)
                    c = read(io, Char)#; print(c)
                    if c == '\'' && !eof(io)
                        c = read(io, Char)#; print(c)
                        if c == ' '
                            if !eof(io)
                                c = read(io, Char)#; print(c)
                            end
                        elseif c == '\n'
                            1 # NO-OP
                        else
                            inchunk = false
                            if !eof(io)
                                c = read(io, Char)#; print(c)
                            end
                        end
                    else
                        inchunk = false
                        if !eof(io)
                            c = read(io, Char)#; print(c)
                        end
                    end
                else
                    inchunk = false
                    if !eof(io)
                        c = read(io, Char)#; print(c)
                    end
                end
            end
        end
    end

    close(io)

    return chunks
end


text:
#=
This is just a copy of Compologicircuits.jl, the code is copy and pasted from that, except the struct Augmented_Circuit has
a new field called sol, the solutions of the circuit, or the graph of the boolean function implied by the circuits
sol is intended to have the type Vector{Vector{BitVector}}, as a vector of pairs [input, output], ordered as
binary numbers on the inputs in increasing order, for instance if f.circ.dom = 5 then
f.sol[1][1] = [0,0,0,0,0], f.sol[2][1] = [0,0,0,0,1], f.sol[3][1] = [0,0,0,1,0], f.sol.[31][1] = [1,1,1,1,1]
Augmented_Circuit might be a useful structure
=#
include("Circuits.jl")
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Theories
using Catlab.Syntax
using Catlab.GAT
using Catlab.Present
using Catlab.Programs
using Catlab.Graphics
using Catlab.Graphics.Graphviz
using Catlab.Programs
using Catlab.WiringDiagrams
using ..Compologicircuits: CircuitDom, Circuit, iNOT, iAND, iOR, Impl, Circuits, show_diagram
import Catlab.Theories: id, compose, otimes, braid, munit, mcopy, delete, pair, proj1, proj2, ⊗, ⋅, σ

B = Ob(FreeCartesianCategory, :B)
dup = mcopy(B)
I = id(B)
and = Hom(:AND, B⊗B, B)
or = Hom(:OR, B⊗B, B)
xor = Hom(:XOR, B⊗B,B)
not = Hom(:NOT, B,B)

to_diagram(x) = add_junctions!(to_wiring_diagram(x))

normal_form(x) = to_hom_expr(FreeCartesianCategory, to_diagram(x))

#returns set of all identical pairs of n length bool vectors
function all_match(n::Int)
    output = []
    for i in 0:(2^n)-1
        vect = .!Vector{Bool}(iszero.(digits(i, base=2, pad=n)))
        push!(output, [vect, vect])
    end
    return output
end

function all_inputs(n::Int)
    output = Vector{Vector{Bool}}()
    for i in 0:(2^n)-1
        vect = .!Vector{Bool}(iszero.(digits(i, base=2, pad=n)))
        push!(output, vect)
    end
    return output
end

#binary vectors to int
function bivec_int(w)::Int
    output = 0
    v = Vector{Int}(w)

    for i in 0:length(v)-1
        output += (v[i+1])*2^i
    end

    return output
end

struct Augmented_CircuitDom
    n::Int
end

struct Augmented_Circuit
    circ::Circuit
    sol::Vector{Vector{Vector{Bool}}} #array of type Vector{Vector{Vector{Bool}} matching input to output
end

Augmented_Circuit(c::Circuit) = begin
    n_ins = c.dom.n
    inputs = all_inputs(n_ins)
    truth_table = Vector{Vector{Vector{Bool}}}()
    for input in inputs
        output = c.impl(input)
        push!(truth_table, [input, output])
    end

    Augmented_Circuit(c, truth_table)
end

aiNOT = Augmented_Circuit(iNOT, [ [[false],[true]], [[true],[false]]])

aiAND = Augmented_Circuit(iAND, [ [[false, false],[false]], [[false, true],[false]], [[true, false],[false]], [[true, true],[true]]  ])

aiOR = Augmented_Circuit(iOR, [ [[false, false],[false]], [[false, true],[true]], [[true, false],[true]], [[true, true],[true]]  ])

#=@instance CartesianCategory{Augmented_CircuitDom, Augmented_Circuit} begin
    id(A::Augmented_CircuitDom) = Augmented_Circuit(Circuit(A,A, x->x), all_match(A.n))
    dom(f::Augmented_Circuit) = f.circ.dom
    codom(f::Augmented_Circuit) = f.circ.codom

    # Circuit composition is just function composition
    compose(f::Augmented_Circuit, g::Augmented_Circuit) = begin
        @assert(f.circ.codom == g.circ.dom)
        output = []
        for i in f.sol
            push!(output, [i[1], g.sol[bivec_int(i[2])+1][2]])
        end
        return Augmented_Circuit(Circuit(f.circ.dom, g.circ.codom, x->g.impl(f.impl(x))), output)
    end

    otimes(A::Augmented_CircuitDom, B::Augmented_CircuitDom) = Augmented_CircuitDom(A.n + B.n)
    # Monoidal product of circuits runs both circuits (TODO: run them in parallel to improve performance) and concatenates the results
    otimes(f::Augmented_Circuit, g::Augmented_Circuit) = begin
        impl = x -> vcat(f.circ.impl(x[1:f.circ.dom.n]), g.circ.impl(x[f.circ.dom.n + 1:end]))
        output = []
        for i in g.sol
            for j in f.sol
                push!(output, [vcat(j[1],i[1]), vcat(j[2],i[2])] )
            end
        end
        return Augmented_Circuit(Circuit(f.circ.dom.n + g.circ.dom.n, f.circ.codom.n + g.circ.codom.n, impl), output)
    end

    # A symmetric braiding just swaps the A and B parts of the input vector
    braid(A::Augmented_CircuitDom, B::Augmented_CircuitDom) = begin
        impl = x -> vcat(x[A.n+1:end], x[1:A.n])
        n = A.n + B.n
        output = all_match(n)
        output = map(v->[v[1],vcat(v[2][A.n+1:end], v[2][1:A.n])], output)
        Augmented_Circuit(Circuit(n, n, impl), output)
    end

    # Monoidal unit is a 0-dimensional bool space (i.e. a point or empty list)
    munit(::Type{Augmented_CircuitDom}) = Augmented_CircuitDom(0)

    # Stuff for Cartesian Category
    mcopy(A::Augmented_CircuitDom) = begin
        output = all_match(A.n)
        output = map(v->[v[1],vcat(v[1],v[1])], output)
        return Augmented_Circuit(Circuit(A.n, A.n+A.n, x->vcat(x,x)), output)
    end

    delete(A::Augmented_CircuitDom) = Augmented_Circuit(Circuit(A.n, 0, x->Bool[]), all_match(0))

    pair(f::Augmented_Circuit, g::Augmented_Circuit) = begin
        @assert(f.circ.dom == g.circ.dom)
        output = []
        for i in f.sol
            push!(output, [i[1], vcat(i[2], g.sol[bivec_int(i[1])+1][2])])
        end
        return Augmented_Circuit(
            Circuit(f.cric.dom.n, f.circ.codom.n+g.circ.codom.n, x->vcat(f.circ.impl(x), g.circ.impl(x))), output)

    end
    proj1(A::Augmented_CircuitDom, B::Augmented_CircuitDom) = begin
        output = all_match(A.n+B.n)
        output = map(v->[v[1], v[2][1:A.n]], output)
        return Augmented_Circuit(Circuit(A.n+B.n, A.n, x->x[1:A.n]), output)
    end
    proj2(A::Augmented_CircuitDom, B::Augmented_CircuitDom) = begin
        output = all_match(A.n+B.n)
        output = map(v->[v[1], v[2][A.n+1:end]], output)
        return Augmented_Circuit(Circuit(A.n+B.n, B.n, x->x[A.n+1:end]), output)
    end
end

gens = Dict(
    Circuits[:B]=>Augmented_CircuitDom(1),
    Circuits[:NOT]=>aiNOT,
    Circuits[:AND]=>aiAND,
    Circuits[:OR]=>aiOR
)

AImpl(expr) = functor((Augmented_CircuitDom, Augmented_Circuit), expr, generators=gens)=#


text:
export Descent, Momentum, AdaGrad, Adam

const ε = 1e-8

struct Descent
    eta::Float64
end

Descent() = Descent(0.1)

apply!(opt::Descent, x, Δ) = opt.eta .* Δ

struct Momentum
    eta::Float64
    alpha::Float64
    velocity::Dict
end

Momentum(η = 0.01, α = 0.9) = Momentum(η, α, Dict())

function apply!(opt::Momentum, x, Δ)
    η, α = opt.eta, opt.alpha
    v = get!(() -> zero(x), opt.velocity, x)::typeof(x) #get! () function returns a view.
    v = (α*v - η*Δ)
    Δ = -v
end

struct AdaGrad
    eta::Float64
    h::Dict
end

AdaGrad(eta = 0.01) = AdaGrad(eta, Dict())

function apply!(opt::AdaGrad, x, Δ)
    η = opt.eta
    h = get!(() -> zero(x), opt.h, x)::typeof(x)
    @. h += Δ^2
    @. Δ *= η / (√h + ε)
end

struct Adam
    eta::Float64
    beta::Tuple{Float64, Float64}
    recode::Dict
end

Adam(η = 0.01, β = (0.9, 0.99)) = Adam(η, β, Dict())

function apply!(opt::Adam, x, Δ)
    η, β = opt.eta, opt.beta
    mt, vt , βp = get!(opt.recode, x) do
        (zero(x), zero(x), Float64[β[1], β[2]])
    end
    @. mt = β[1] * mt + (1 - β[1]) * Δ
    @. vt = β[2] * vt + (1 - β[2]) * Δ^2
    @. Δ = η * mt / (√(vt / (1 - βp[2])) + ε)
    βp .= βp .* β
    return Δ
end


text:
function solve()
    N = parse(Int64, readline())
    S = [chomp(readline()) for s in 1:N]

	numMARCH = [0, 0, 0, 0, 0]
    MARCH = ["M", "A", "R", "C", "H"]

  	for name in S
        for (i, m) in enumerate(MARCH)
            if name[1] == m[1]
                numMARCH[i] += 1
            end
        end
    end

    nonzero = numMARCH[numMARCH .!= 0]
    if length(nonzero) < 3
        println(0)
        return
    end

  	combinations::Int64 = 0
    for i in 1:length(nonzero)
    	for j in i+1:length(nonzero)
            for k in j+1:length(nonzero)
                combinations += nonzero[i]*nonzero[j]*nonzero[k]
            end
        end
    end
    println(combinations)

end

solve()


text:
"""
    f_stats(x::PopObj, mode::String = "global")
Returns a DataFrame of the F-statistics of a `PopObj`. See
`FST` to perform pairwise FST comparisons between populations.

### mode options
- `global` : provides FIT, FIS, and FST for an entire population (default)
- `loci` : provides FIS, FIT, and FST for each locus
- `sample` : provides FIS and FIT for each sample
"""
function f_stats(x::PopObj, mode::String = "global")

    if mode == "locus" || mode == "loci"

        fst_locus(x)

    elseif mode == "sample" || mode == "ind" || mode == "individual"

        fst_sample(x)

    elseif mode == "global"

        fst_global(x)

    else
        error("please specify \"global\", \"loci\", or \"sample\" as the second argument")
    end
end

"""
    fst_global(data::PopObj, round::Bool = true)
Calculates the three common F-statistics Fᵢₛ, Fᵢₜ, and Fₛₜ on a PopObj, adjusted for
different sample sizes in populations. If `round` is `true` (default), the results are
rounded to 4 decimal places, otherwise `false` returns the entire values as calculated.
"""
function fst_global(data::PopObj, round::Bool = true)

    ## H_i = (Hobs1*N1 + Hobs2*N2 + etc..) / Ntotal
    popcounts = populations(data).count
    pop_hets = het(data, "pop")
    obs = map(i -> mean(i.het_obs), pop_hets) |> collect
    exp = map(i -> mean(i.het_exp), pop_hets) |> collect
    H_i = sum(obs .* popcounts) / sum(popcounts)

    ## n/n-1 correction
    correction = sum(popcounts) / (sum(popcounts) - 1)
    ## H_s = (Hexp1*N1 + Hexp2*N2 + etc..) / Ntotal
    H_s = (sum(exp .* popcounts) / sum(popcounts)) * correction

    ## H_t = 1 - ∑(p²+q²)     (expected Het)
    H_t = mean(het(data).het_exp)

    @info "Global F-statistics"

    # FIS = (HS - HI) / HS
    FIS = (H_s - H_i) / H_s

    # FIT = (HT- HI)/HT
    FIT = (H_t - H_i) / H_t

    # FST = (HT- HS) / HT
    FST = (H_t - H_s) / H_t

    if round == false
        return DataFrame(
            FIS = FIS,
            FIT = FIT,
            FST = FST
        )
    else
        return DataFrame(
            FIS = Base.round.(FIS, digits = 4),
            FIT = Base.round.(FIT, digits = 4),
            FST = Base.round.(FST, digits = 4),
        )
    end
end

function fst_locus(data::PopObj)
    hets = He(data, "pop")
    H_t = het_expected(data) |> mean
    H_s = map(i -> mean(i.het_exp), hets)
    H_i = map(i -> mean(i.het_obs), hets)

    @info "F-statistics by Locus"

    # FIT = (HT- HI)/HT
    FIT = map(i -> (H_t - i) / H_t, het_observed(data))

    # FIS = (HS - HI) / HS
    HI = het_observed(data)
    HS = map(i -> i.het_exp, hets)
    FIS = []
    for (i, j) in enumerate(HI)
        locus = map(a -> a[i], HS)
        fis = (locus .- j) ./ locus |> mean
        push!(FIS, fis)
    end
    FIS = round.(FIS, digits = 4) |> Array{Float64,1}

    # FST = (HT- HS) / HT
    FST = []
    for i = 1:length(loci(data))
        locus = map(c -> c[i], HS)
        fst = (H_t .- locus) ./ H_t |> mean
        push!(FST, fst)
    end
    FST = round.(FST, digits = 4) |> Array{Float64,1}

    return DataFrame(
        locus = loci(x),
        FIS = FIS,
        FIT = round.(FIT, digits = 4),
        FST = FST,
    )
end

function fst_sample(data::PopObj)
    H_i_sample = het_sample(x)
    @info "F-statistics by Sample"
    # split sample heterozygosities into their own GroupedDataFrames
    insertcols!(H_i_sample, 2, :population => x.samples.population)
    H_i_sample_pop = groupby(H_i_sample[!, :], :population)

    FIS = []
    for (i, j) in enumerate(H_i_sample_pop)
        # FIS = (HS - HI) / HS
        tmp = map(k -> (H_s[i] - k) / H_s[i], j.het)
        append!(FIS, tmp)
    end

    FIS = round.(FIS, digits = 4) |> Array{Float64,1}

    # FIT = (HT- HI)/HT
    # H_t =
    FIT = map(i -> (H_t - i) / H_t, H_i_sample.het)

    return DataFrame(
        name = samples(x),
        FIS = FIS,
        FIT = round.(FIT, digits = 4),
    )
end

"""
    f_stat_p(data::PopObj, nperm::Int64)
Performs a permutation test (permuting individuals among populations) to calculate
p-values for F-statistics
"""
function f_stat_p(data::PopObj, nperm::Int64 = 1000)
    tmp = deepcopy(data)

    observed = f_stats(data, "global")
    sims = deepcopy(observed)
    for boot in 1:(nperm-1)
        shuffle!(tmp.samples.population)
        append!(sims, f_stats(tmp, "global"))
    end
    return map(i -> sum((observed .<= sims)[!,i])/nperm, 1:ncol(observed))
end


text:
"""
    Imputor

An imputor stores information about imputing values in `AbstractArray`s and `Tables.table`s.
New imputation methods are expected to sutype `Imputor` and, at minimum,
implement the `impute!(imp::<MyImputor>, data::AbstractVector)` method.
"""
abstract type Imputor end

# A couple utility methods to avoid messing up var and obs dimensions
obsdim(dims) = dims
vardim(dims) = dims == 1 ? 2 : 1

function obswise(data::AbstractMatrix; dims=1)
    return (selectdim(data, obsdim(dims), i) for i in axes(data, obsdim(dims)))
end

function varwise(data::AbstractMatrix; dims=2)
    return (selectdim(data, vardim(dims), i) for i in axes(data, vardim(dims)))
end

function filterobs(f::Function, data::AbstractMatrix; dims=1)
    mask = [f(x) for x in obswise(data; dims=dims)]
    return dims == 1 ? data[mask, :] : data[:, mask]
end

function filtervars(f::Function, data::AbstractMatrix; dims=2)
    mask = [f(x) for x in varwise(data; dims=dims)]
    return dims == 1 ? data[:, mask] : data[mask, :]
end

"""
    splitkwargs(::Type{T}, kwargs...) where T <: Imputor -> (imp, rem)

Takes an Imputor type with kwargs and returns the constructed imputor and the
unused kwargs which should be passed to the `impute!` call.

NOTE: This is used by utility methods with construct and imputor and call impute in 1 call.
"""
function splitkwargs(::Type{T}, kwargs...) where T <: Imputor
    rem = Dict(kwargs...)
    kwdef = empty(rem)

    for f in fieldnames(T)
        if haskey(rem, f)
            kwdef[f] = rem[f]
            delete!(rem, f)
        end
    end

    return (T(; kwdef...), rem)
end

# Some utility methods for constructing imputors and imputing data in 1 call.
# NOTE: This is only intended for internal use and is not part of the public API.
function _impute(data, t::Type{T}, kwargs...) where T <: Imputor
    imp, rem = splitkwargs(t, _extract_context_kwargs(kwargs...)...)
    return impute(data, imp; rem...)
end

function _impute!(data, t::Type{T}, kwargs...) where T <: Imputor
    imp, rem = splitkwargs(t, _extract_context_kwargs(kwargs...)...)
    return impute!(data, imp; rem...)
end

"""
    impute(data, imp::Imputor; kwargs...)

Returns a new copy of the `data` with the missing data imputed by the imputor `imp`.

# Keywords
* `dims`: The dimension to impute along (e.g., observations dim)
"""
function impute(data, imp::Imputor; kwargs...)
    # Call `deepcopy` because we can trust that it's available for all types.
    return impute!(deepcopy(data), imp; kwargs...)
end

"""
    impute!(data::AbstractMatrix, imp::Imputor; kwargs...)

Impute the data in a matrix by imputing the values one variable at a time;
if this is not the desired behaviour custom imputor methods should overload this method.

# Arguments
* `data::AbstractMatrix`: the data to impute
* `imp::Imputor`: the Imputor method to use

# Keywords
* `dims`: The dimension to impute along (e.g., observations dim)

# Returns
* `AbstractMatrix`: the input `data` with values imputed

# Example
```jldoctest
julia> using Impute: Interpolate, Context, impute

julia> M = [1.0 2.0 missing missing 5.0; 1.1 2.2 3.3 missing 5.5]
2×5 Array{Union{Missing, Float64},2}:
 1.0  2.0   missing  missing  5.0
 1.1  2.2  3.3       missing  5.5

julia> impute(M, Interpolate(; context=Context(; limit=1.0)); dims=2)
2×5 Array{Union{Missing, Float64},2}:
 1.0  2.0  3.0  4.0  5.0
 1.1  2.2  3.3  4.4  5.5
```
"""
function impute!(data::AbstractMatrix, imp::Imputor; dims=1)
    for var in varwise(data; dims=dims)
        impute!(var, imp)
    end
    return data
end

"""
    impute!(table, imp::Imputor)

Imputes the data in a table by imputing the values 1 column at a time;
if this is not the desired behaviour custom imputor methods should overload this method.

# Arguments
* `imp::Imputor`: the Imputor method to use
* `table`: the data to impute

# Returns
* the input `data` with values imputed

# Example
``jldoctest
julia> using DataFrames; using Impute: Interpolate, Context, impute
julia> df = DataFrame(:a => [1.0, 2.0, missing, missing, 5.0], :b => [1.1, 2.2, 3.3, missing, 5.5])
5×2 DataFrame
│ Row │ a        │ b        │
│     │ Float64⍰ │ Float64⍰ │
├─────┼──────────┼──────────┤
│ 1   │ 1.0      │ 1.1      │
│ 2   │ 2.0      │ 2.2      │
│ 3   │ missing  │ 3.3      │
│ 4   │ missing  │ missing  │
│ 5   │ 5.0      │ 5.5      │

julia> impute(df, Interpolate(; context=Context(; limit=1.0)))
5×2 DataFrame
│ Row │ a        │ b        │
│     │ Float64⍰ │ Float64⍰ │
├─────┼──────────┼──────────┤
│ 1   │ 1.0      │ 1.1      │
│ 2   │ 2.0      │ 2.2      │
│ 3   │ 3.0      │ 3.3      │
│ 4   │ 4.0      │ 4.4      │
│ 5   │ 5.0      │ 5.5      │
"""
function impute!(table, imp::Imputor)
    istable(table) || throw(MethodError(impute!, (table, imp)))

    # Extract a columns iterator that we should be able to use to mutate the data.
    # NOTE: Mutation is not guaranteed for all table types, but it avoid copying the data
    columntable = Tables.columns(table)

    for cname in propertynames(columntable)
        impute!(getproperty(columntable, cname), imp)
    end

    return table
end

for file in ("drop.jl", "locf.jl", "nocb.jl", "interp.jl", "fill.jl", "chain.jl", "srs.jl")
    include(joinpath("imputors", file))
end


text:
"""
`MIRTio`
I/O routines for Michigan Image Reconstruction Toolbox (MIRT)
"""
module MIRTio

include("fastmri/hdf5utils.jl")

include("ge-mri/header.jl")
include("ge-mri/pfile.jl")
include("ge-mri/rdb-26_002.jl")
include("ge-mri/read_rdb_hdr.jl")

end # module


text:
function bisection_search(f::Function, x0::Vector, x1::Vector, t::Vector;
        tol::Real = 1e-5,
        maxiters::Int = 1000,
        verbose = true
    )

    @assert all(x0 .<= x1)
    @assert all(length.([x0, x1, t]) .== length(x0))

    x0, x1 = (x0, x1) .|> deepcopy

    # finding sense
    r0 = f(x0)
    r1 = f(x1)
    @assert length(r0) == length(r1) == length(x0)
    sense = map((x) -> x[1] <= x[2] ? 1.0 : -1.0, zip(r0, r1))
    for (i, (tc, sc, rc0, rc1)) in enumerate(zip(t, sense, r0, r1))
        !(sc*rc0 <= tc <= sc*rc1) && error("target[$i] = $tc not in range: [$(sc*rc0), $(sc*rc1)]")
    end

    verbose && (prog = ProgressThresh(tol, "searching:  "))
    xi = x1
    last_xi = x1
    for it in 1:maxiters
        xi = (x1 - x0)/2 + x0
        ri = f(xi)

        # Assume a constant monotony TODO: rethink this
        foreach(eachindex(x0)) do c
            sense[c] > 0.0 ?
                ri[c] < t[c] ? (x0[c] = xi[c]) : (x1[c] = xi[c]) :
                ri[c] < t[c] ? (x1[c] = xi[c]) : (x0[c] = xi[c])
        end

        error = maximum(abs.(ri - t))
        verbose && update!(prog, error)
        if all(error < tol)
            verbose && finish!(prog)
            return xi
        end
        last_xi = xi
    end
    verbose && finish!(prog)
    return xi
end


text:
@field ShadingFunction = blinnphong

@composed type MeshUniforms
    Model
    Color
    ShadingFunction
end

function vert_mesh(vertex::AbstractVertex, canvas, light, shading, uniforms)
    #args = (solid_color, proj, view)
    pos = to_vec4(getposition(vertex))
    mv = canvas.view * uniforms.model
    position_camspace = mv * pos
    # direction to light
    lightdir = normalize(light.position .- pos[Vec(1, 2, 3)])
    # direction to camera
    # screen space coordinates of the vertex
    pos_screen = canvas.projection * position_camspace
    return (
        pos_screen,
        Vert2Frag(
            getcolor(vertex, uniforms),
            -position_camspace[Vec(1,2,3)],
            getnormal(vertex),
            getuv(vertex),
            lightdir
        )
    )
end

function frag_mesh(vertex_out::Vert2Frag, canvas, light, shading, uniforms)
    # (solid_color, proj, view)
    color = vertex_out.color
    V = vertex_out.position
    L = normalize(vertex_out.lightdir)
    N = normalize(vertex_out.normal)

    c = uniforms.shadingfunction(V, N, L, color[Vec(1, 2, 3)], shading, light)
    (Vec4f0(c[1], c[2], c[3], color[4]),)
end

@composed type Mesh
    <: MeshUniforms
    Vertices
    Light
    Shading
end

function Drawable(window::AbstractWindow, mesh::Mesh)
    shading = mesh[Shading]
    # Create a functor
    vertices = mesh[Vertices]

    uniforms = (window[Scene], window[Light], shading, MeshUniforms(mesh))

    rasterizer(
        window,
        vertices, uniforms,
        vert_mesh, Visualize.frag_mesh
    )
end


text:
@safetestset "REM" begin
    using GlobalMatchingModels, Statistics, Test, Random
    using Distributions
    probe = [6,1,1,3];
    memory = [0 2; 1 2; 0 1; 3 0];
    model = REM(;memory, g=.40, c=.70)
    activations = compute_activations(model, probe)
    @test activations ≈ [10.5802,.18450] atol = 1e-4

    prob = compute_prob(activations)
    @test prob ≈ .8433 atol = 1e-4

    Random.seed!(58)
    model = REM(;memory=fill(0, 4, 4), g=.40, c=.70)
    stimuli = generate_stimuli(.3, 4, 4)
    encode!(model, stimuli)
    dims = size(model.memory)
    @test dims == (4,4)
    @test sum(model.memory .!= stimuli) > 0

    Random.seed!(25)
    g = .30
    c = .70
    u = .80
    n = 5000
    stimuli = generate_stimuli(g, n, n)
    model = REM(;memory=fill(0, n, n), g, c, u)
    encode!(model, stimuli)
    est = mean(model.memory .== 0)
    @test est ≈  (1 - u) + u * g * c + u * (1 - c) * g atol = 1e-4
end


text:
using BlockArrays
using DataFrames
using LinearAlgebra
using MixedModels
using NamedArrays
using Random
using SparseArrays
using Statistics
using Tables
using Test

const LMM = LinearMixedModel
const forms = Dict(
    :dyestuff => [@formula(yield ~ 1 + (1|batch))],
    :dyestuff2 => [@formula(yield ~ 1 + (1|batch))],
    :d3 => [@formula(y ~ 1 + u + (1+u|g) + (1+u|h) + (1+u|i))],
    :insteval => [
        @formula(y ~ 1 + service + (1|s) + (1|d) + (1|dept)),
        @formula(y ~ 1 + service*dept + (1|s) + (1|d)),
    ],
    :kb07 => [
        @formula(rt_trunc ~ 1+spkr+prec+load+(1|subj)+(1|item)),
        @formula(rt_trunc ~ 1+spkr*prec*load+(1|subj)+(1+prec|item)),
        @formula(rt_trunc ~ 1+spkr*prec*load+(1+spkr+prec+load|subj)+(1+spkr+prec+load|item)),
    ],
    :ml1m => [@formula(Y ~ 1 + (1|G) + (1|H))],
    :pastes => [
        @formula(strength ~ 1 + (1|sample)),
        @formula(strength ~ 1 + (1|sample) + (1|batch)),
    ],
    :penicillin => [@formula(diameter ~ 1 + (1|plate) + (1|sample))],
    :sleepstudy => [
        @formula(reaction ~ 1 + days + (1|subj)),
        @formula(reaction ~ 1 + days + zerocorr(1+days|subj)),
        @formula(reaction ~ 1 + days + (1+days|subj)),
    ],
)

#=         # replace all the model fits by
fits = Dict(k => fit.(MixedModel, forms[k], Ref(MixedModels.dataset[k])) for k in keys(forms))
=#

@testset "Dyestuff" begin
    ds = MixedModels.dataset(:dyestuff)
    fm1 = LMM(@formula(yield ~ 1 + (1|batch)), ds)

    @test length(fm1.allterms) == 3
    @test size(fm1.reterms) == (1, )
    @test lowerbd(fm1) == zeros(1)
    @test fm1.lowerbd == zeros(1)
    @test fm1.θ == ones(1)
    fm1.θ = ones(1)
    @test fm1.θ == ones(1)
    @test_logs (:warn, "Model has not been fit") show(fm1)

    @test objective(updateL!(setθ!(fm1, [0.713]))) ≈ 327.34216280955366
    @test_deprecated MixedModels.describeblocks(IOBuffer(), fm1)

    io = IOBuffer()
    show(io, BlockDescription(fm1))
    @test countlines(seekstart(io)) == 3
    output = String(take!(io))
    @test startswith(output, "rows:")

    fit!(fm1);
    @test :θ in propertynames(fm1)
    @test objective(fm1) ≈ 327.3270598811428 atol=0.001
    @test fm1.θ ≈ [0.752580] atol=1.e-5
    @test fm1.λ ≈ [LowerTriangular(reshape(fm1.θ, (1,1)))] atol=1.e-5
    @test deviance(fm1) ≈ 327.32705988 atol=0.001
    @test aic(fm1) ≈ 333.3270598811394 atol=0.001
    @test bic(fm1) ≈ 337.5306520261259 atol=0.001
    @test fixef(fm1) ≈ [1527.5]
    @test first(first(fm1.σs)) ≈ 37.26034462135931 atol=0.0001
    @test fm1.β ≈ [1527.5]
    @test dof(fm1) == 3
    @test nobs(fm1) == 30
    @test MixedModels.fixef!(zeros(1),fm1) ≈ [1527.5]
    @test coef(fm1) ≈ [1527.5]
    fm1β = fm1.βs
    @test fm1β isa NamedTuple
    @test isone(length(fm1β))
    @test first(values(fm1β)) ≈ 1527.5
    fm1σρ = fm1.σρs
    @test fm1σρ isa NamedTuple
    @test isone(length(fm1σρ))
    @test isone(length(getproperty(first(fm1σρ), :σ)))
    @test isempty(getproperty(first(fm1σρ), :ρ))
    @test fm1.σ == sdest(fm1)
    @test fm1.b == ranef(fm1)
    @test fm1.u == ranef(fm1, uscale=true)
    @test fm1.stderror == stderror(fm1)
    @test isone(length(fm1.pvalues))
    @test fm1.objective == objective(fm1)
    @test fm1.σ ≈ 49.510099986291145 atol=1.e-5
    @test fm1.X == ones(30,1)
    @test fm1.y == ds[!, :yield]
    @test cond(fm1) == ones(1)
    @test first(leverage(fm1)) ≈ 0.15650534392640486 rtol=1.e-5
    @test sum(leverage(fm1)) ≈ 4.695160317792145 rtol=1.e-5
    cm = coeftable(fm1)
    @test length(cm.rownms) == 1
    @test length(cm.colnms) == 4
    @test fnames(fm1) == (:batch,)
    @test response(fm1) == ds[!, :yield]
    rfu = ranef(fm1, uscale = true)
    rfb = ranef(fm1)
    @test abs(sum(rfu[1])) < 1.e-5
    cv = condVar(fm1)
    @test length(cv) == 1
    @test size(first(cv)) == (1, 1, 6)
    show(IOBuffer(), fm1.optsum)

    @test logdet(fm1) ≈ 8.06014522999825 atol=0.001
    @test varest(fm1) ≈ 2451.2501089607676 atol=0.001
    @test pwrss(fm1) ≈ 73537.49947885796 atol=0.001
    @test stderror(fm1) ≈ [17.69455188898009] atol=0.0001

    vc = VarCorr(fm1)
    show(IOBuffer(), vc)
    @test vc.s == sdest(fm1)

    fit!(fm1, REML=true)
    @test objective(fm1) ≈ 319.65427684225216 atol=0.0001
    @test_throws ArgumentError loglikelihood(fm1)
    @test dof_residual(fm1) ≥ 0
    print(IOBuffer(), fm1)
end

@testset "Dyestuff2" begin
    ds2 = MixedModels.dataset(:dyestuff2)
    fm = fit(MixedModel, @formula(yield ~ 1 + (1|batch)), ds2)
    @test lowerbd(fm) == zeros(1)
    show(IOBuffer(), fm)
    @test fm.θ ≈ zeros(1)
    @test objective(fm) ≈ 162.87303665382575
    @test abs(std(fm)[1][1]) < 1.0e-9
    @test std(fm)[2] ≈ [3.653231351374652]
    @test stderror(fm) ≈ [0.6669857396443261]
    @test coef(fm) ≈ [5.6656]
    @test logdet(fm) ≈ 0.0
    refit!(fm, float(MixedModels.dataset(:dyestuff)[!, :yield]))
    @test objective(fm) ≈ 327.3270598811428 atol=0.001
end

@testset "penicillin" begin
    pen = MixedModels.dataset(:penicillin)
    fm = LMM(@formula(diameter ~ 1 + (1 | plate) + (1 | sample)), pen);
    @test size(fm) == (144, 1, 30, 2)
    @test fm.θ == ones(2)
    @test lowerbd(fm) == zeros(2)

    fit!(fm)

    @test objective(fm) ≈ 332.18834867227616 atol=0.001
    @test coef(fm) ≈ [22.97222222222222] atol=0.001
    @test fixef(fm) ≈ [22.97222222222222] atol=0.001
    @test coef(fm)[1] ≈ mean(pen.diameter)
    @test stderror(fm) ≈ [0.7445960346851368] atol=0.0001
    @test fm.θ ≈ [1.5375772376554968, 3.219751321180035] atol=0.001
    @test first(std(fm)) ≈ [0.8455645948223015] atol=0.0001
    @test std(fm)[2] ≈ [1.770647779277388] atol=0.0001
    @test varest(fm) ≈ 0.3024263987592062 atol=0.0001
    @test logdet(fm) ≈ 95.74614821367786 atol=0.001
    rfu = ranef(fm, uscale=true)
    rfb = ranef(fm)
    @test length(rfb) == 2
end

@testset "pastes" begin
    fm = LMM(@formula(strength ~ (1|sample) + (1|batch)), MixedModels.dataset(:pastes))
    @test size(fm) == (60, 1, 40, 2)
    @test fm.θ == ones(2)
    @test lowerbd(fm) == zeros(2)

    fit!(fm);

    @test objective(fm) ≈ 247.99446586289676 atol=0.001
    @test coef(fm) ≈ [60.05333333333329] atol=0.001
    @test fixef(fm) ≈ [60.05333333333329] atol=0.001
    @test stderror(fm) ≈ [0.6421359883527029] atol=0.0001
    @test fm.θ ≈ [3.5268858714382905, 1.3299230213750168] atol=0.001
    @test first(std(fm)) ≈ [2.904069002535747] atol=0.001
    @test std(fm)[2] ≈ [1.095070371687089] atol=0.0001
    @test std(fm)[3] ≈ [0.8234088395243269] atol=0.0001
    @test varest(fm) ≈ 0.6780020742644107 atol=0.0001
    @test logdet(fm) ≈ 101.0381339953986 atol=0.001
end

@testset "InstEval" begin
    insteval = MixedModels.dataset(:insteval)
    fm1 = LMM(@formula(y ~ 1 + service + (1|s) + (1|d) + (1|dept)), insteval)
    @test size(fm1) == (73421, 2, 4114, 3)
    @test fm1.θ == ones(3)
    @test lowerbd(fm1) == zeros(3)

    fit!(fm1);

    @test objective(fm1) ≈ 237721.7687745563 atol=0.001
    ftd1 = fitted(fm1);
    @test size(ftd1) == (73421, )
    @test ftd1 == predict(fm1)
    @test first(ftd1) ≈ 3.17876 atol=0.0001
    resid1 = residuals(fm1);
    @test size(resid1) == (73421, )
    @test first(resid1) ≈ 1.82124 atol=0.00001

    @testset "PCA" begin
        @test length(fm1.rePCA) == 3
        @test length(MixedModels.PCA(fm1)) == 3
    end

    fm2 = fit(MixedModel, @formula(y ~ 1 + service*dept + (1|s) + (1|d)), insteval)
    @test objective(fm2) ≈ 237585.5534151694 atol=0.001
    @test size(fm2) == (73421, 28, 4100, 2)
end

@testset "sleep" begin
    slp = MixedModels.dataset(:sleepstudy)
    fm = LinearMixedModel(@formula(reaction ~ 1 + days + (1+days|subj)), slp);
    @test lowerbd(fm) == [0.0, -Inf, 0.0]
    A11 = fm.A[Block(1,1)]
    @test isa(A11, UniformBlockDiagonal{Float64})
    @test isa(fm.L[Block(1, 1)], UniformBlockDiagonal{Float64})
    @test size(A11) == (36, 36)
    @test A11.facevec[1] == [10. 45.; 45. 285.]
    @test length(A11.facevec) == 18
    updateL!(fm);
    b11 = LowerTriangular(fm.L[Block(1, 1)].facevec[1])
    @test b11 * b11' == fm.A[Block(1, 1)].facevec[1] + I
    @test count(!iszero, Matrix(fm.L[Block(1, 1)])) == 18 * 4
    @test rank(fm) == 2

    fit!(fm)

    @test objective(fm) ≈ 1751.9393444647046
    @test fm.θ ≈ [0.929221307, 0.01816838, 0.22264487096] atol=1.e-6
    @test pwrss(fm) ≈ 117889.46144025437
    @test logdet(fm) ≈ 73.90322021999222 atol=0.001
    @test stderror(fm) ≈ [6.632257721914501, 1.5022354739749826] atol=0.0001
    @test coef(fm) ≈ [251.40510484848477,10.4672859595959]
    @test fixef(fm) ≈  [251.40510484848477,10.4672859595959]
    @test std(fm)[1] ≈ [23.780468100188497, 5.716827903196682] atol=0.01
    @test logdet(fm) ≈ 73.90337187545992 atol=0.001
    @test cond(fm) ≈ [4.175251] atol=0.0001
    @test loglikelihood(fm) ≈ -875.9696722323523
    @test sum(leverage(fm)) ≈ 28.611525700136877 rtol=1.e-5
    σs = fm.σs
    @test length(σs) == 1
    @test keys(σs) == (:subj,)
    @test length(σs.subj) == 2
    @test first(values(σs.subj)) ≈ 23.780468626896745 atol=0.0001
    @test last(values(first(σs))) ≈ 5.716827808126002 atol=0.0001
    show(IOBuffer(), fm)

    @test fm.vcov ≈ [43.986843599069985 -1.370391879315539; -1.370391879315539 2.2567113314443645] atol=0.0001
    @test fm.corr ≈ [1.0 -0.1375451787621904; -0.1375451787621904 1.0] atol=0.0001

    u3 = ranef(fm, uscale=true)
    @test length(u3) == 1
    @test size(first(u3)) == (2, 18)
    @test first(u3)[1, 1] ≈ 3.030300122575336 atol=0.001
    u3n = first(ranef(fm, uscale=true, named=true))
    @test u3n isa NamedArrays.NamedArray

    b3 = ranef(fm)
    @test length(b3) == 1
    @test size(first(b3)) == (2, 18)
    @test first(first(b3)) ≈ 2.815819441982976 atol=0.001

    simulate!(fm)  # to test one of the unscaledre methods

    fmnc = zerocorr!(deepcopy(fm))
    @test fmnc.optsum.feval < 0
    @test size(fmnc) == (180,2,36,1)
    @test fmnc.θ == [fm.θ[1], fm.θ[3]]
    @test lowerbd(fmnc) == zeros(2)

    fmnc = zerocorr!(LMM(@formula(reaction ~ 1 + days + (1+days|subj)), slp));
    @test size(fmnc) == (180,2,36,1)
    @test fmnc.θ == ones(2)
    @test lowerbd(fmnc) == zeros(2)

    @testset "zerocorr PCA" begin
        @test length(fmnc.rePCA) == 1
        @test fmnc.rePCA.subj == [0.5, 1.0]
        @test MixedModels.PCA(fmnc).subj.loadings == I(2)
        @test show(IOBuffer(), MixedModels.PCA(fmnc)) === nothing
    end

    fit!(fmnc)

    @test deviance(fmnc) ≈ 1752.0032551398835 atol=0.001
    @test objective(fmnc) ≈ 1752.0032551398835 atol=0.001
    @test coef(fmnc) ≈ [251.40510484848585, 10.467285959595715]
    @test fixef(fmnc) ≈ [251.40510484848477, 10.467285959595715]
    @test stderror(fmnc) ≈ [6.707710260366577, 1.5193083237479683] atol=0.001
    @test fmnc.θ ≈ [0.9458106880922268, 0.22692826607677266] atol=0.0001
    @test first(std(fmnc)) ≈ [24.171449463289047, 5.799379721123582]
    @test last(std(fmnc)) ≈ [25.556130034081047]
    @test logdet(fmnc) ≈ 74.46952585564611 atol=0.001
    ρ = first(fmnc.σρs.subj.ρ)
    @test ρ === -0.0   # test that systematic zero correlations are returned as -0.0

    fmnc2 = LinearMixedModel(@formula(reaction ~ 1+days+zerocorr(1+days|subj)), slp)
    @test size(fmnc2) == (180,2,36,1)
    @test fmnc2.θ == ones(2)
    @test lowerbd(fmnc2) == zeros(2)

    fit!(fmnc2)

    @test deviance(fmnc2) ≈ 1752.0032551398835 atol=0.001
    @test objective(fmnc2) ≈ 1752.0032551398835 atol=0.001
    @test coef(fmnc2) ≈ [251.40510484848585, 10.467285959595715]
    @test fixef(fmnc2) ≈ [251.40510484848477, 10.467285959595715]
    @test stderror(fmnc2) ≈ [6.707710260366577, 1.5193083237479683] atol=0.001
    @test fmnc2.θ ≈ [0.9458106880922268, 0.22692826607677266] atol=0.0001
    @test first(std(fmnc2)) ≈ [24.171449463289047, 5.799379721123582]
    @test last(std(fmnc2)) ≈ [25.556130034081047]
    @test logdet(fmnc2) ≈ 74.46952585564611 atol=0.001

    MixedModels.likelihoodratiotest(fm, fmnc)

    fmrs = fit(MixedModel, @formula(reaction ~ 1+days + (0+days|subj)), slp);
    @test objective(fmrs) ≈ 1774.080315280528 rtol=0.00001
    @test fmrs.θ ≈ [0.24353985679033105] rtol=0.00001

    fm_ind = fit(MixedModel, @formula(reaction ~ 1+days + (1|subj) + (0+days|subj)), slp)
    @test objective(fm_ind) ≈ objective(fmnc)
    @test coef(fm_ind) ≈ coef(fmnc)
    @test fixef(fm_ind) ≈ fixef(fmnc)
    @test stderror(fm_ind) ≈ stderror(fmnc)
    @test fm_ind.θ ≈ fmnc.θ
    @test std(fm_ind) ≈ std(fmnc)
    @test logdet(fm_ind) ≈ logdet(fmnc)

    # combining [ReMat{T,S1}, ReMat{T,S2}] for S1 ≠ S2
    slpcat = categorical!(deepcopy(slp), [:days])
    fm_cat = LMM(@formula(reaction ~ 1 + days + (1|subj) + (0 + days|subj)), slpcat)
    @test fm_cat isa LMM
    σρ = fit!(fm_cat).σρs
    @test σρ isa NamedTuple
    @test isone(length(σρ))
    @test first(keys(σρ)) == :subj
    @test keys(σρ.subj) == (:σ, :ρ)
    @test length(σρ.subj) == 2
    @test length(first(σρ.subj)) == 11
    @test length(σρ.subj.ρ) == 55
    @test iszero(σρ.subj.ρ[46])
    @test σρ.subj.ρ[46] === -0.0
end

@testset "d3" begin
    fm = updateL!(LMM(@formula(y ~ 1 + u + (1+u|g) + (1+u|h) + (1+u|i)), MixedModels.dataset(:d3)));
    @test pwrss(fm) ≈ 5.1261847180180885e6 rtol = 1e-6
    @test objective(fm) ≈ 901641.2930413672 rtol = 1e-6
    fit!(fm)
    @test objective(fm) ≈ 884957.5540213 rtol = 1e-6
    @test coef(fm) ≈ [0.4991229873, 0.31130780953] atol = 1.e-4
    bv = ranef(fm)
    @test length(bv) == 3
end

@testset "simulate!" begin
    ds = MixedModels.dataset(:dyestuff)
    fm = fit(MixedModel, @formula(yield ~ 1 + (1|batch)), ds)
    resp₀ = copy(response(fm))
    # type conversion of ints to floats
    simulate!(Random.MersenneTwister(1234321), fm, β=[1], σ=1)
    refit!(fm,resp₀)
    refit!(simulate!(Random.MersenneTwister(1234321), fm))
    @test deviance(fm) ≈ 339.0218639362958 atol=0.001
    refit!(fm, float(ds.yield))
    Random.seed!(1234321)
    refit!(simulate!(fm))
    @test deviance(fm) ≈ 339.0218639362958 atol=0.001
    simulate!(fm, θ = fm.θ)
    @test_throws DimensionMismatch refit!(fm, zeros(29))

    # type conversion of ints to floats
    parametricbootstrap(Random.MersenneTwister(1234321), 1, fm, β=[1], σ=1)
    # this should give the same results as passing with
    # MersenneTwister(1234321)
    # by setting the seed first, we test the method for the default RNG
    # and still get reproducible results which can be compared to the
    # multi-threaded output.
    Random.seed!(1234321)
    bsamp = parametricbootstrap(100, fm, use_threads=false)
    @test isa(propertynames(bsamp), Vector{Symbol})
    @test length(bsamp.objective) == 100
    @test keys(first(bsamp.bstr)) == (:objective, :σ, :β, :se, :θ)
    @test isa(bsamp.σs, Vector{<:NamedTuple})
    @test length(bsamp.σs) == 100
    allpars = DataFrame(bsamp.allpars)
    @test isa(allpars, DataFrame)
    cov = shortestcovint(bsamp.σ)
    @test first(cov) ≈ 48.2551828768727 rtol = 1.e-4
    @test last(cov) ≈ 81.85810781858969 rtol = 1.e-4

    bsamp_threaded = parametricbootstrap(MersenneTwister(1234321), 100, fm, use_threads=true)
    # even though it's bad practice with floating point, exact equality should
    # be a valid test here -- if everything is working right, then it's the exact
    # same operations occuring within each bootstrap sample, which IEEE 754
    # guarantees will yield the same result
    @test sort(bsamp_threaded.σ) == sort(bsamp.σ)
    @test sort(bsamp_threaded.θ) == sort(bsamp.θ)
    @test sort(columntable(bsamp_threaded.β).β) == sort(columntable(bsamp.β).β)
end

@testset "Rank deficient" begin
    rng = MersenneTwister(0);
    x = rand(rng, 100);
    data = columntable((x = x, x2 = 1.5 .* x, y = rand(rng, 100), z = repeat('A':'T', 5)))
    model = fit(MixedModel, @formula(y ~ x + x2 + (1|z)), data)
    @test length(fixef(model)) == 2
    @test rank(model) == 2
    @test length(coef(model)) == 3
    ct = coeftable(model)
    @test ct.rownms ==  ["(Intercept)", "x", "x2"]
    @test fixefnames(model) == ["(Intercept)", "x"]
    @test coefnames(model) == ["(Intercept)", "x", "x2"]
    @test last(coef(model)) == -0.0
    stde = MixedModels.stderror!(zeros(3), model)
    @test isnan(last(stde))

    # check preserving of name ordering in coeftable and placement of
    # pivoted-out element
    fill!(data.x, 0)
    model2 = fit(MixedModel, @formula(y ~ x + x2 + (1|z)), data)
    ct = coeftable(model2)
    @test ct.rownms ==  ["(Intercept)", "x", "x2"]
    @test fixefnames(model2) == ["(Intercept)", "x2"]
    @test coefnames(model2) == ["(Intercept)", "x", "x2"]
    @test coef(model2)[2] == -0.0
    @test last(fixef(model)) ≈ (last(fixef(model2)) * 1.5)
    stde2 = MixedModels.stderror!(zeros(3), model2)
    @test isnan(stde2[2])

end

@testset "coeftable" begin
    ds = MixedModels.dataset(:dyestuff);
    fm = fit(MixedModel, @formula(yield ~ 1 + (1|batch)), ds);
    ct = coeftable(fm);
    @test [3,4] == [ct.teststatcol, ct.pvalcol]
end


text:
using FizzBuzzQuiz
using Test

tests = ["fizz_buzz", "pezz", "hozz"]
for test = tests
    # println("testing $test.jl...")
    include("$test.jl")
end


text:
using Test
using ProxAL
using DelimitedFiles, Printf
using LinearAlgebra, JuMP
using CatViews
using CUDA

DATA_DIR = joinpath(dirname(@__FILE__), "..", "data")
case = "case9"
T = 2
ramp_scale = 0.5
load_scale = 1.0
rtol = 1e-4

# Load case
case_file = joinpath(DATA_DIR, "$(case).m")
load_file = joinpath(DATA_DIR, "mp_demand", "$(case)_oneweek_168")

# Model/formulation settings
modelinfo = ModelInfo()
modelinfo.num_time_periods = T
modelinfo.load_scale = load_scale
modelinfo.ramp_scale = ramp_scale
modelinfo.allow_obj_gencost = true
modelinfo.allow_constr_infeas = false
modelinfo.weight_freq_ctrl = 0.1
modelinfo.time_link_constr_type = :penalty

# Algorithm settings
algparams = AlgParams()
algparams.verbose = 0

solver_list = ["Ipopt", "ExaTron"]
# TODO: MadNLP broken currently
# solver_list = ["Ipopt", "MadNLP"]
# if CUDA.has_cuda_gpu()
#     push!(solver_list, "MadNLPGPU")
# end
if isfile(joinpath(dirname(@__FILE__), "..", "build/libhiop.so"))
    push!(solver_list, "Hiop")
    ENV["JULIA_HIOP_LIBRARY_PATH"] = joinpath(dirname(@__FILE__), "..", "build")
    @info("Using Hiop at $(ENV["JULIA_HIOP_LIBRARY_PATH"])")
end

@testset "Test ProxAL on $(case)" begin
    modelinfo.case_name = case

    for solver in solver_list
    @testset "$(solver)" begin
        println("Testing using $(solver)")
        if solver == "Ipopt"
            using Ipopt
            backend = JuMPBackend()
            algparams.optimizer =
                optimizer_with_attributes(Ipopt.Optimizer,
                    "print_level" => Int64(algparams.verbose > 0)*5)
        end
        if solver == "ExaTron"
            using ExaTron
            backend = ExaTronBackend()
            algparams.tron_outer_iterlim=2000
            algparams.tron_outer_eps=1e-6
        end

        @testset "solver = $(solver)" begin

            K = 0
            algparams.decompCtgs = false
            @testset "$T-period, $K-ctgs, time_link=$(modelinfo.time_link_constr_type)" begin
                modelinfo.num_ctgs = K
                OPTIMAL_OBJVALUE = round(11258.316096599736*modelinfo.obj_scale, digits = 6)
                OPTIMAL_PG = round.([0.8979870694509675, 1.3432060120295906, 0.9418738103137331, 0.9840203268625166, 1.448040098924617, 1.0149638876964715], digits = 5)

                if solver == "Ipopt"
                    @testset "Non-decomposed formulation" begin
                        algparams.mode = :nondecomposed
                        algparams.θ_t = algparams.θ_c = (1/algparams.tol)^2
                        nlp = NonDecomposedModel(case_file, load_file, modelinfo, algparams)
                        result = ProxAL.optimize!(nlp)
                        @test isapprox(result["objective_value_nondecomposed"], OPTIMAL_OBJVALUE, rtol = rtol)
                        @test isapprox(result["primal"].Pg[:], OPTIMAL_PG, rtol = rtol)
                        @test norm(result["primal"].Zt[:], Inf) <= algparams.tol
                    end
                end

                @testset "ProxAL" begin
                    algparams.mode = :coldstart
                    nlp = ProxALEvaluator(case_file, load_file, modelinfo, algparams, backend, Dict(), Dict(), nothing)
                    runinfo = ProxAL.optimize!(nlp)
                    @test isapprox(runinfo.objvalue[end], OPTIMAL_OBJVALUE, rtol = rtol)
                    @test isapprox(runinfo.x.Pg[:], OPTIMAL_PG, rtol = rtol)
                    @test isapprox(runinfo.maxviol_c[end], 0.0)
                    @test isapprox(runinfo.maxviol_c_actual[end], 0.0)
                    @test runinfo.maxviol_t[end] <= algparams.tol
                    @test runinfo.maxviol_t_actual[end] <= algparams.tol
                    @test runinfo.maxviol_d[end] <= algparams.tol
                    @test runinfo.iter <= algparams.iterlim
                end
            end

            K = 1
            modelinfo.num_ctgs = K
            for ctgs_link in [:frequency_equality, :preventive_equality, :corrective_inequality]
                if ctgs_link == :frequency_equality
                    OPTIMAL_OBJVALUE = round(11258.316096601551*modelinfo.obj_scale, digits = 6)
                    OPTIMAL_PG = round.([0.8979870771416967, 1.3432060071608862, 0.9418738040358301, 0.9840203279072699, 1.448040097565594, 1.0149638851897345], digits = 5)
                    OPTIMAL_WT = round.([0.0, -0.0001206958510371199, 0.0, -0.0001468516901269684], sigdigits = 4)
                    proxal_ctgs_link = :frequency_penalty
                elseif ctgs_link == :preventive_equality
                    OPTIMAL_OBJVALUE = round(11340.093581315918*modelinfo.obj_scale, digits = 6)
                    OPTIMAL_PG = round.([0.8802072707975164, 1.3551904372529155, 0.9621162845248026, 0.9632463589831285, 1.46270646294457, 1.03946802455576], digits = 5)
                    OPTIMAL_WT = [0.0, 0.0, 0.0, 0.0]
                    proxal_ctgs_link = :preventive_penalty
                elseif ctgs_link == :corrective_inequality
                    OPTIMAL_OBJVALUE = round(11258.316096599581*modelinfo.obj_scale, digits = 6)
                    OPTIMAL_PG = round.([0.8979870771538515, 1.3432060071606775, 0.941873804023475, 0.9840203279201325, 1.4480400975658037, 1.0149638851762735], digits = 5)
                    OPTIMAL_WT = [0.0, 0.0, 0.0, 0.0]
                    proxal_ctgs_link = :corrective_penalty
                end
                if solver == "ExaTron" && proxal_ctgs_link != :corrective_penalty
                    continue
                end

                algparams.decompCtgs = false
                modelinfo.ctgs_link_constr_type = ctgs_link
                if solver == "Ipopt"
                    @testset "$T-period, $K-ctgs, time_link=$(modelinfo.time_link_constr_type), ctgs_link=$(ctgs_link)" begin
                        @testset "Non-decomposed formulation" begin
                            algparams.mode = :nondecomposed
                            algparams.θ_t = algparams.θ_c = (1/algparams.tol)^2
                            nlp = NonDecomposedModel(case_file, load_file, modelinfo, algparams)
                            result = ProxAL.optimize!(nlp)
                            @test isapprox(result["objective_value_nondecomposed"], OPTIMAL_OBJVALUE, rtol = rtol)
                            @test isapprox(result["primal"].Pg[:,1,:][:], OPTIMAL_PG, rtol = rtol)
                            @test norm(result["primal"].Zt[:], Inf) <= algparams.tol
                            @test isapprox(result["primal"].ωt[:], OPTIMAL_WT, rtol = 1e-1)
                        end

                        @testset "ProxAL" begin
                            algparams.mode = :coldstart
                            nlp = ProxALEvaluator(case_file, load_file, modelinfo, algparams, backend, Dict(), Dict(), nothing)
                            runinfo = ProxAL.optimize!(nlp)
                            @test isapprox(runinfo.objvalue[end], OPTIMAL_OBJVALUE, rtol = rtol)
                            @test isapprox(runinfo.x.Pg[:,1,:][:], OPTIMAL_PG, rtol = rtol)
                            @test isapprox(runinfo.x.ωt[:], OPTIMAL_WT, rtol = 1e-1)
                            @test isapprox(runinfo.maxviol_c[end], 0.0)
                            @test isapprox(runinfo.maxviol_c_actual[end], 0.0)
                            @test runinfo.maxviol_t[end] <= algparams.tol
                            @test runinfo.maxviol_t_actual[end] <= algparams.tol
                            @test runinfo.maxviol_d[end] <= algparams.tol
                            @test runinfo.iter <= algparams.iterlim
                        end
                    end
                end

                algparams.decompCtgs = true
                modelinfo.ctgs_link_constr_type = proxal_ctgs_link
                @testset "$T-period, $K-ctgs, time_link=$(modelinfo.time_link_constr_type), ctgs_link=$(proxal_ctgs_link), decompCtgs" begin
                    if solver == "Ipopt"
                        @testset "Non-decomposed formulation" begin
                            algparams.mode = :nondecomposed
                            algparams.θ_t = algparams.θ_c = (10/algparams.tol)^2
                            nlp = NonDecomposedModel(case_file, load_file, modelinfo, algparams)
                            result = ProxAL.optimize!(nlp)
                            @test isapprox(result["objective_value_nondecomposed"], OPTIMAL_OBJVALUE, rtol = rtol)
                            @test isapprox(result["primal"].Pg[:,1,:][:], OPTIMAL_PG, rtol = rtol)
                            @test norm(result["primal"].Zt[:], Inf) <= algparams.tol
                            @test norm(result["primal"].Zk[:], Inf) <= algparams.tol
                        end
                    end

                    @testset "ProxAL" begin
                        algparams.mode = :coldstart
                        algparams.iterlim = 300
                        nlp = ProxALEvaluator(case_file, load_file, modelinfo, algparams, backend, Dict(), Dict(), nothing)
                        runinfo = ProxAL.optimize!(nlp)
                        @test isapprox(runinfo.objvalue[end], OPTIMAL_OBJVALUE, rtol = 1e-2)
                        @test isapprox(runinfo.x.Pg[:,1,:][:], OPTIMAL_PG, rtol = 1e-2)
                        @test runinfo.maxviol_c[end] <= algparams.tol
                        @test runinfo.maxviol_t[end] <= algparams.tol
                        @test runinfo.maxviol_c_actual[end] <= algparams.tol
                        @test runinfo.maxviol_t_actual[end] <= algparams.tol
                        @test runinfo.maxviol_d[end] <= algparams.tol
                        @test runinfo.iter <= algparams.iterlim
                    end
                end
            end
        end
    end # solver testset
    end
end