text:
##################################################################################
# This file is part of StateSpaceEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

"""
    presolve_sstate!(model; <options>)
    presolve_sstate!(model, mask, values; <options>)

Solve for the steady state variables that are decoupled from the
system, or can be solved by forward substitution.

This is called automatically by the steady state solver before
running its main loop.

### Arguments
  * `model` - the Model instance
  * `mask` - a vector of Bool. Defaults to `model.sstate.mask`
  * `values` - a vector of numbers. Defaults to `model.sstate.values`
Caller must specify either both `mask` and `values` or neither of them.
`mask[i]` equals `true` if and only if the i-th steady state value has
alredy been solved for.
`mask` and `values` are both input and output data to the algorithm.
Any values that are successfully presolved are updated in place,
and their `mask` enties are set to `true`.

### Options
  * `verbose` - `true` or `false`, whether or not to print diagnostic messages.
  * `tol` - accuracy of the 1d solver.

"""
@inline presolve_sstate!(model::Model; kwargs...) = presolve_sstate!(model, model.sstate.mask, model.sstate.values; kwargs...)
function presolve_sstate!(model::Model, mask::AbstractVector{Bool}, values::AbstractVector{Float64};
    verbose = model.options.verbose, tol = model.options.tol,
    _1dsolver = :bisect)

    if _1dsolver == :newton
        solve1d = (sseqn, vals, ind)->newton1!(sseqn.eval_RJ, vals, ind; maxiter = 5, tol = tol)
    elseif _1dsolver == :bisect
        solve1d = (sseqn, vals, ind)->begin
            _, J = sseqn.eval_RJ(vals)
            bisect!(sseqn.eval_resid, vals, ind, J[ind]; maxiter = 1000, tol = tol)
        end
    else
        error("Unknown 1d solver $(_1dsolver)")
    end
    local alleqns = ModelBaseEcon.alleqns(model.sstate)
    while true
        # Keep looping for as long as we keep updating values
        updated = false
        for (eqind, sseqn) ∈ enumerate(alleqns)
            unsolved = .! mask[sseqn.vinds]
            if sum(unsolved) == 1
                # There is exactly one variable in this equation that's not solved yet.
                ind = findall(unsolved)[1]
                vals = view(values, sseqn.vinds)
                sym = ModelBaseEcon.ss_symbol(model.sstate, sseqn.vinds[ind])
                # store the current value, so we can restore it if 1d solver fails.
                foo = vals[ind]

                if solve1d(sseqn, vals, ind)
                    if abs(vals[ind]) < 1e-10
                        vals[ind] = 0.0
                    end
                    if verbose
                        pt_info = NamedTuple{tuple(sseqn.vsyms...)}(vals)
                        @info "Presolved $eqind: $sseqn for $sym = $(vals[ind])" pt_info
                    end
                    mask[sseqn.vinds[ind]] = true
                    updated = true
                else
                    vals[ind] = foo
                    if verbose
                        pt_info = NamedTuple{tuple(sseqn.vsyms...)}(vals)
                        @info "Failed to presolve $eqind: $sseqn for $sym" pt_info
                    end
                end
            end
        end
        #
        if ! updated
            break
        end
    end
    nothing
end
@assert precompile(presolve_sstate!, (Model, Vector{Bool}, Vector{Float64}))


text:
using YaoSym, YaoBlocks, YaoArrayRegister, SymEngine
using YaoSym: simplify_expi
using Test

@testset "mat" begin
    @vars θ γ η
    for G in [X, Y, Z, ConstGate.T, H]
        @test Matrix(mat(Basic, G)) ≈ Matrix(mat(G))
        @test Matrix(mat(Basic, control(4, 3, 2=>G))) ≈ Matrix(control(4, 3, 2=>G))
    end

    for GC in [Rx, shift, phase]
        G = GC(θ)
        m = mat(Basic, G)
        m = subs.(m, θ, 0.5)
        @test Matrix(mat(GC(0.5))) ≈ Matrix(m)
        @test subs.(Matrix(mat(Basic, put(4, 2=>G))), θ, 0.5) ≈ Matrix(put(4, 2=>GC(0.5)))
    end

    G = Rz(θ)*Rx(γ)*Rz(θ)
    m = mat(Basic, G)
    @test_throws ArgumentError mat(Float64, G)
    m = subs.(m, Ref(θ=>Basic(π)/2), Ref(γ=>Basic(π)/6))
    @test Matrix(mat(Rz(π/2)*Rx(π/6)*Rz(π/2))) ≈ Matrix(m)

    A = randn(ComplexF64, 4,4)
    mb = matblock(Basic.(A))
    @test Matrix(mb) ≈ A

    @test pswap(4,2,1,θ) == pswap(2,1,θ)(4)
    @test pswap(4,2,1,θ) isa PSwap{4,Basic}
end

ext: jl


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

using Markdown
using InteractiveUtils

# ╔═╡ efff10ac-824b-422b-b2f9-4074ea3437fb
begin
	using Pkg; Pkg.activate(@__DIR__); Pkg.instantiate()
	Pkg.precompile()

	using PlutoUI
	using UMAP
	using Makie
	using GLMakie
	using XLSX
	using VegaDatasets
	using DataFrames
	using MultivariateStats
	using RDatasets
	using StatsBase
	using Statistics
	using LinearAlgebra
	using Plots
	using ScikitLearn
	using MLBase
	using Distances
end;

# ╔═╡ 85e2d832-542c-49b0-a26d-b57a7a3d4404
PlutoUI.TableOfContents(aside=true, indent=true, depth=3)

# ╔═╡ 6e8c97d2-4d9a-4f00-99e1-69054903635c
html"""

"""

# ╔═╡ 133df7ea-d3ec-41fa-bcb8-5eae4c853a55
md"""
# 📝 Dimensionality Reduction

As the name says, dimensionality reduction is the idea of reducing your feature set to a much smaller number. Dimensionality reduction is often used in visualization of datasets to try and detect samples that are similar. Usually, the goal is to come up with a 2D visualization of your data, where every dot in your plot is representing an observation. We can use these plots to look for patterns in data. We will cover three dimensionality reduction techniques here:

1. t-SNE
2. PCA
3. UMAP

[YouTube Link](https://www.youtube.com/watch?v=hIsYy04zO7U)
"""

# ╔═╡ 5271ffe0-f73a-430e-8ccc-5c6bd0937ce0
md"""
## Load dataset

We will use a dataset from the [VegaDatasets.jl](https://github.com/queryverse/VegaDatasets.jl) package. The dataset is about car specifications of over 400 car models.
"""

# ╔═╡ 6a53752f-3101-4b62-a661-c09c9df27f0f
data = VegaDatasets.dataset("cars") |> DataFrame

# ╔═╡ 9a491a1d-32a0-496c-887a-cdfd49f01836
md"""
Let's check the available variables with the `names` function...
"""

# ╔═╡ dcc6939b-ba8e-4f4a-abcf-aa90efa60682
names(data)

# ╔═╡ 56529677-8b5c-4912-b25a-728847b4f374
md"""
## Data wrangling

"""

# ╔═╡ 59d179aa-a19c-497d-ae56-6e0b0b0007dc
md"""
We will first drop missing data with the `dropmissing!` inplace function...
"""

# ╔═╡ 2f0ba288-19b9-49c8-845c-bdc2317328fc
dropmissing!(data)

# ╔═╡ 182b795f-c069-4ecd-aab6-89c66c93ee6e
md"""
Now, we need to save the numeric variables (columns 2 through 7)  as a `Matrix`. We will use this data to perform dimensionality reduction methods later...
"""

# ╔═╡ 56e3357f-a692-4b29-8ca2-390ba62b9679
M = data[:,2:7] |> Matrix

# ╔═╡ 53bae393-698c-4610-b18f-073a67d9bdc8
md"""
We are also going to use the `Origin` feature. However, since it is a vector of strings, we will need to encode it.

First, we will use the [MLBase.jl](https://github.com/JuliaStats/MLBase.jl) `labelmap` function to map `Origin` labels and, then, we will encode it using the `labelencode` function...
"""

# ╔═╡ 193d7109-5499-40c3-a25f-b7b8560db5b9
origin = data[!,:Origin]

# ╔═╡ c3d8e866-c754-4ee6-8c1c-fab7f8048181
origin_map = labelmap(origin)

# ╔═╡ 231955b0-5ccd-46ad-b78d-96f861f1938a
origin_codes = labelencode(origin_map, origin)

# ╔═╡ 841ef716-8dbf-4aec-947e-6c0409a7fd64
md"""
**Note:** after encoding, we have:
- USA → 1
- Japan → 2
- Europe → 3
"""

# ╔═╡ 2c86f61a-b328-4eac-a4ff-1e713007ef54
md"""
Finally, we must center and normalize the data...

**Note:** `dims=1` means that we are calculating both `mean` and `std` over _columns_.
"""

# ╔═╡ c08882ca-d617-4c0f-865f-898c0f639a82
center_M = (M .- mean(M, dims=1)) ./ std(M, dims=1)

# ╔═╡ a9afbe75-29a4-4a54-b6c2-2786d74f7a65
md"""
## PCA

**Principal component analysis (PCA)** is the process of computing the principal components and using them to perform a change of basis on the data, sometimes using only the first few principal components and ignoring the rest.

The principal components of a collection of points in a real coordinate space are a sequence of $p$ unit vectors, where the $i$-th vector is the direction of a line that best fits the data while being orthogonal to the first $i-1$ vectors. Here, a best-fitting line is defined as one that minimizes the average squared distance from the points to the line. These directions constitute an orthonormal basis in which different individual dimensions of the data are linearly uncorrelated.

**Note:** the details for the underlying mathematics can be found [here](https://en.wikipedia.org/wiki/Principal_component_analysis).

First, we will fit the model via PCA, using the [MultivariateStats.jl](https://github.com/JuliaStats/MultivariateStats.jl) package. `maxoutdim` is the output dimensions and we want it to be 2 in this case.

**Note:** since PCA expects each column to be an observation, we will use the transpose of the matrix.
"""

# ╔═╡ dffd6373-2e05-4746-bc44-ee2193ae2635
pca = fit(PCA, center_M', maxoutdim=2)

# ╔═╡ d13d825e-e807-4ef0-a9ab-45fc994500a2
md"""
Now, we can obtain the projection matrix by calling the function `projection`...

**Note:** once you multiply the projection matrix by your data, you get the dimensionality reduction of your data.
"""

# ╔═╡ cb2f6b28-d360-4bb8-a0c0-3d6d1b38e23d
P = projection(pca)

# ╔═╡ 8d89f524-68a0-425e-a4e9-2525ebdd914d
md"""
Given a PCA model, one can use it to transform observations into principal components, as:

```math
y = P^T(x - \mu)
```

where $P$ is the projection matrix, $x$ is the original feature vector and $\mu$  is the empirical mean vector by each original feature. We can perform this transformation manually or call `transform` function...
"""

# ╔═╡ 3946ec51-0336-4098-ba13-55abf68f9b41
# empirical mean vector
μ = mean(pca)

# ╔═╡ 8d219680-c870-4549-9421-5c7fff3a7607
# transforming the 1st instance manually
P' * (center_M[1,:] - μ)

# ╔═╡ a450f277-46cf-4f5b-b58a-b33b5aca5fb2
md"""
**Note:** the result above is the 1st feature vector represented by the first two principal components.
"""

# ╔═╡ 2e2d14e9-8fdb-4328-9218-f47307309f40
# transforming all data
Y = MultivariateStats.transform(pca, center_M')

# ╔═╡ 13b52e6a-dbe4-4dad-a7a6-260957fb0bdd
md"""
**Note:** `Y[:,1]` is the same as `P' * (center_data[1,:] - μ)`:
"""

# ╔═╡ 49f05ab7-b716-44c4-bb44-6c2e700a18ac
Y[:,1] ≈ (P' * (center_M[1,:] - μ))

# ╔═╡ c37484f0-3e35-4c5f-841a-e0e6061752b4
md"""
Given a PCA model, we also may use it to reconstruct (approximately) the observations from principal components, as:

```math
\tilde{x} = Py + \mu
```
where $P$ is the projection matrix, $x$ is the original feature vector and $\mu$  is the empirical mean vector by each original feature. We can perform this reconstruction using the `reconstruct` function...
"""

# ╔═╡ cf067613-8312-402c-88b2-fe507ce36ac9
X̅ = reconstruct(pca, Y)

# ╔═╡ 7594fb9c-443a-4ea8-a56a-f17c818ddbfe
md"""
However, since the reconstruction is approximate, the norm difference between $\tilde{x}$ and $x$ will not be zero...
"""

# ╔═╡ e0e09b17-da85-4879-bec8-adef9f125b31
norm(X̅ - center_M')

# ╔═╡ 0e73cd39-24d3-4de4-a30d-a53ad6f1d50c
md"""
Finally, we can generate a scatter plot of the cars...
"""

# ╔═╡ 5dfc1915-f832-4d72-b551-e3ddbfa693bc
Plots.scatter(Y[1,:], Y[2,:], color=:orange, legend=false,
			  xlabel="1st PC", ylabel="2nd PC")

# ╔═╡ 98c5d1d2-6c28-4c46-adeb-bb580d5dfe2a
md"""
Let's try to color the scatter plot by `Origin`...
"""

# ╔═╡ 8117e9fa-23d1-4539-aa28-c5bc4a3db849
begin
	# boolean lists
	usaᵢ = origin.=="USA"
	japᵢ = origin .== "Japan"
	eurᵢ = origin .== "Europe"
end;

# ╔═╡ 07dbb182-4565-48de-b41c-bd0488f67280
begin
	Plots.scatter(Y[1,usaᵢ], Y[2,usaᵢ], color=1, label="USA")
	Plots.scatter!(Y[1,japᵢ], Y[2,japᵢ], color=2, label="Japan")
	Plots.scatter!(Y[1,eurᵢ], Y[2,eurᵢ], color=3, label="Europe")

	Plots.xlabel!("1st PC")
	Plots.ylabel!("2nd PC")
end

# ╔═╡ 6bca80ac-8c1d-4f35-82e7-4c6c4206eb53
md"""
This is interesting! There seems to be three main clusters with cars from the US dominating two clusters.

Now, let's try to reduce the dimentionality to 3 principal components...
"""

# ╔═╡ 152fc839-661b-48aa-a047-9e16a35acff5
begin
	new_pca = fit(PCA, center_M', maxoutdim=3)
	new_Y = MultivariateStats.transform(new_pca, center_M')
end

# ╔═╡ 5f55ba88-c984-49d2-b4b9-e5cdd81e8c6c
scatter3d(new_Y[1,:], new_Y[2,:], new_Y[3,:], color=origin_codes,
		  xlabel="1st PC", ylabel="2nd PC", zlabel="3rd PC",
		  legend=false, camera=(50, 30))

# ╔═╡ bac92bf8-3059-4fcc-b1a9-42c91637fd17
md"""
## t-SNE

**t-Distributed Stochastic Neighbor Embedding (t-SNE)** is a statistical method for visualizing high-dimensional data by giving each datapoint a location in a two or three-dimensional map. It is based on Stochastic Neighbor Embedding originally developed by Sam Roweis and Geoffrey Hinton,[1] where Laurens van der Maaten proposed the t-distributed variant.

**Note:** the details for the underlying mathematics can be found [here](https://en.wikipedia.org/wiki/T-distributed_stochastic_neighbor_embedding).

The next method we will use for dimensionality reduction is t-SNE. There are multiple ways you can call t-SNE from Julia. But we will take this opportunity to try out something new... We will call a function from the scikitlearn Python package. This makes use of the package [ScikitLearn.jl](https://github.com/cstjean/ScikitLearn.jl).

**Note:** check this [t-SNE notebook](https://github.com/nassarhuda/JuliaTutorials/blob/master/TSNE/TSNE.ipynb).
"""

# ╔═╡ 6c163758-55b5-474b-8269-203884d70757
# importing sklearn t-SNE
@sk_import manifold: TSNE

# ╔═╡ 4676988f-5687-4aad-a054-48ae5defd46e
# instantiate a t-SNE object with 2 components
tsne = TSNE(n_components=2, perplexity=20.0, early_exaggeration=50)

# ╔═╡ 026668de-5323-4e40-b75c-c2eefbf46ef5
# fit-transforming data
Y₂ = tsne.fit_transform(center_M)

# ╔═╡ 7fb412b7-573e-4c5a-bd39-d9d7e637a9d4
md"""
**Note:** here, we do not have to work with the transpose data, since ScikitLearn already deals with observations as rows...

Let's plot the components...
"""

# ╔═╡ 8faa2556-bcfc-4289-9d17-d75ad10f2245
begin
	Plots.scatter(Y₂[usaᵢ,1], Y₂[usaᵢ,2], color=1, label="USA")
	Plots.scatter!(Y₂[japᵢ,1], Y₂[japᵢ,2], color=2, label="Japan")
	Plots.scatter!(Y₂[eurᵢ,1], Y₂[eurᵢ,2], color=3, label="Europe")

	Plots.xlabel!("1st Component")
	Plots.ylabel!("2nd Component")
end

# ╔═╡ 3114d2e5-e6eb-4198-98f3-4abfbb45bdeb
md"""
This is interesting! The same patterns saw before appears to hold here too.
"""

# ╔═╡ 7675ae48-a6ae-4189-990f-c1469ee11929
md"""
## UMAP

**Uniform Manifold Approximation and Projection (UMAP)** is a dimension reduction technique that can be used for visualisation similarly to t-SNE, but also for general non-linear dimension reduction. The algorithm is founded on three assumptions about the data:

1. The data is uniformly distributed on Riemannian manifold;
2. The Riemannian metric is locally constant (or can be approximated as such);
3. The manifold is locally connected.

From these assumptions it is possible to model the manifold with a fuzzy topological structure. The embedding is found by searching for a low dimensional projection of the data that has the closest possible equivalent fuzzy topological structure.

**Note:** the details for the underlying mathematics can be found [here](https://arxiv.org/abs/1802.03426).

To perform UMAP, we will use the [UMAP.jl](https://github.com/dillondaudert/UMAP.jl) package. First, let's calculate the correlation matrix between cars...

**Note:** `dims=2` means that correlation is calculated over rows/instances.
"""

# ╔═╡ bdaff924-889a-4dfc-93cc-4375f6016d82
L = cor(center_M, center_M, dims=2)

# ╔═╡ 3b892bcb-4c67-4309-be2c-5d562e4b98b5
md"""
Now, we can create an UMAP instance with two components...
"""

# ╔═╡ 5f56f6f7-f905-4271-a111-ba4092e3d12e
embed₁ = umap(L, 2)

# ╔═╡ 0cdb33ae-e168-42a3-974a-2a8cf0ccea8c
md"""
**Note:** the output components are being treated as rows (not columns).
"""

# ╔═╡ d70792c1-c3e4-48cb-8168-34cfa89ca0f4
begin
	Plots.scatter(embed₁[1,usaᵢ], embed₁[2,usaᵢ], color=1, label="USA")
	Plots.scatter!(embed₁[1,japᵢ], embed₁[2,japᵢ], color=2, label="Japan")
	Plots.scatter!(embed₁[1,eurᵢ], embed₁[2,eurᵢ], color=3, label="Europe")

	Plots.xlabel!("1st Component")
	Plots.ylabel!("2nd Component")
end

# ╔═╡ 6f34e24a-ef1f-4b33-a639-aec7a8a12aef
md"""
Instead of calculating the correlation matrix, we can calculate the pairwise Euclidean distance between instances using [Distances.jl](https://github.com/JuliaStats/Distances.jl) package...
"""

# ╔═╡ 08a89d70-7f5d-4a5e-8e4a-35abd03bb5b7
D = pairwise(Euclidean(), center_M, center_M, dims=1)

# ╔═╡ 0734fd78-fef5-4536-a99b-57eef9cd5f0f
md"""
Now, we can create again an UMAP instance with two components...
"""

# ╔═╡ d0ed58f2-93c1-4d4e-8f81-5f376c14ec41
embed₂ = umap(D, 2)

# ╔═╡ cef603ec-8e05-43cc-962e-11c763fe09af
begin
	Plots.scatter(embed₂[1,usaᵢ], embed₂[2,usaᵢ], color=1, label="USA")
	Plots.scatter!(embed₂[1,japᵢ], embed₂[2,japᵢ], color=2, label="Japan")
	Plots.scatter!(embed₂[1,eurᵢ], embed₂[2,eurᵢ], color=3, label="Europe")

	Plots.xlabel!("1st Component")
	Plots.ylabel!("2nd Component")
end

# ╔═╡ Cell order:
# ╟─efff10ac-824b-422b-b2f9-4074ea3437fb
# ╟─85e2d832-542c-49b0-a26d-b57a7a3d4404
# ╟─6e8c97d2-4d9a-4f00-99e1-69054903635c
# ╟─133df7ea-d3ec-41fa-bcb8-5eae4c853a55
# ╟─5271ffe0-f73a-430e-8ccc-5c6bd0937ce0
# ╠═6a53752f-3101-4b62-a661-c09c9df27f0f
# ╟─9a491a1d-32a0-496c-887a-cdfd49f01836
# ╠═dcc6939b-ba8e-4f4a-abcf-aa90efa60682
# ╟─56529677-8b5c-4912-b25a-728847b4f374
# ╟─59d179aa-a19c-497d-ae56-6e0b0b0007dc
# ╠═2f0ba288-19b9-49c8-845c-bdc2317328fc
# ╟─182b795f-c069-4ecd-aab6-89c66c93ee6e
# ╠═56e3357f-a692-4b29-8ca2-390ba62b9679
# ╟─53bae393-698c-4610-b18f-073a67d9bdc8
# ╠═193d7109-5499-40c3-a25f-b7b8560db5b9
# ╠═c3d8e866-c754-4ee6-8c1c-fab7f8048181
# ╠═231955b0-5ccd-46ad-b78d-96f861f1938a
# ╟─841ef716-8dbf-4aec-947e-6c0409a7fd64
# ╟─2c86f61a-b328-4eac-a4ff-1e713007ef54
# ╠═c08882ca-d617-4c0f-865f-898c0f639a82
# ╟─a9afbe75-29a4-4a54-b6c2-2786d74f7a65
# ╠═dffd6373-2e05-4746-bc44-ee2193ae2635
# ╟─d13d825e-e807-4ef0-a9ab-45fc994500a2
# ╠═cb2f6b28-d360-4bb8-a0c0-3d6d1b38e23d
# ╟─8d89f524-68a0-425e-a4e9-2525ebdd914d
# ╠═3946ec51-0336-4098-ba13-55abf68f9b41
# ╠═8d219680-c870-4549-9421-5c7fff3a7607
# ╟─a450f277-46cf-4f5b-b58a-b33b5aca5fb2
# ╠═2e2d14e9-8fdb-4328-9218-f47307309f40
# ╟─13b52e6a-dbe4-4dad-a7a6-260957fb0bdd
# ╠═49f05ab7-b716-44c4-bb44-6c2e700a18ac
# ╟─c37484f0-3e35-4c5f-841a-e0e6061752b4
# ╠═cf067613-8312-402c-88b2-fe507ce36ac9
# ╟─7594fb9c-443a-4ea8-a56a-f17c818ddbfe
# ╠═e0e09b17-da85-4879-bec8-adef9f125b31
# ╟─0e73cd39-24d3-4de4-a30d-a53ad6f1d50c
# ╠═5dfc1915-f832-4d72-b551-e3ddbfa693bc
# ╟─98c5d1d2-6c28-4c46-adeb-bb580d5dfe2a
# ╠═8117e9fa-23d1-4539-aa28-c5bc4a3db849
# ╠═07dbb182-4565-48de-b41c-bd0488f67280
# ╟─6bca80ac-8c1d-4f35-82e7-4c6c4206eb53
# ╠═152fc839-661b-48aa-a047-9e16a35acff5
# ╠═5f55ba88-c984-49d2-b4b9-e5cdd81e8c6c
# ╟─bac92bf8-3059-4fcc-b1a9-42c91637fd17
# ╠═6c163758-55b5-474b-8269-203884d70757
# ╠═4676988f-5687-4aad-a054-48ae5defd46e
# ╠═026668de-5323-4e40-b75c-c2eefbf46ef5
# ╟─7fb412b7-573e-4c5a-bd39-d9d7e637a9d4
# ╠═8faa2556-bcfc-4289-9d17-d75ad10f2245
# ╟─3114d2e5-e6eb-4198-98f3-4abfbb45bdeb
# ╟─7675ae48-a6ae-4189-990f-c1469ee11929
# ╠═bdaff924-889a-4dfc-93cc-4375f6016d82
# ╟─3b892bcb-4c67-4309-be2c-5d562e4b98b5
# ╠═5f56f6f7-f905-4271-a111-ba4092e3d12e
# ╟─0cdb33ae-e168-42a3-974a-2a8cf0ccea8c
# ╠═d70792c1-c3e4-48cb-8168-34cfa89ca0f4
# ╟─6f34e24a-ef1f-4b33-a639-aec7a8a12aef
# ╠═08a89d70-7f5d-4a5e-8e4a-35abd03bb5b7
# ╟─0734fd78-fef5-4536-a99b-57eef9cd5f0f
# ╠═d0ed58f2-93c1-4d4e-8f81-5f376c14ec41
# ╠═cef603ec-8e05-43cc-962e-11c763fe09af

max_stars_repo_name: fnaghetini/DataScience

max_issues_repo_name: fnaghetini/DataScience

max_forks_repo_name: fnaghetini/DataScience


text:
# Apply variable selection using power likelihood to the CPP (collaborative perinatal project data)
include("varsel.jl")

module RunCPP
using VarSel
using PyPlot
draw() = (pause(0.001); get_current_fig_manager()[:window][:raise_]())

# Settings
n_total = 10^4  # total number of MCMC samples
n_keep = n_total  # number of MCMC samples to record
alphas = [100,500,1000,2000,Inf]  # values of alpha to use
nburn = int(n_keep/10)  # burn-in

# Variables to use
target_name = "V_BWGT"  # name of target variable
predictor_names = vec(readdlm("datasets/CPP/predictors.txt"))  # names of the predictor variables / covariates

# randomly permute order to check for possible MCMC mixing issues
#predictor_names = predictor_names[randperm(length(predictor_names))]

# Load data
data = readdlm("datasets/CPP/finaldde.csv",',')
names = vec(data[1,:])
values = data[2:end,:]
predictor_subset = [findfirst(names.==name)::Int64 for name in predictor_names]
@assert(minimum(predictor_subset)>0)
target_index = findfirst(names.==target_name)
for j in predictor_subset
    missing = vec(values[:,j].==".")
    values[missing,j] = mean(values[!missing,j])
end
println("Missing values in covariates have been replaced by the average of non-missing entries.")
missing_y = vec(values[:,target_index].==".")
y = convert(Array{Float64,1}, values[!missing_y,target_index])
x = convert(Array{Float64,2}, values[!missing_y,predictor_subset])
n = length(y)
println("Removed $(sum(missing_y)) records which were missing the target variable.")

# Preprocess data
y = (y-mean(y))./std(y)
x = (x.-mean(x,1))./std(x,1)
@assert(all(!isnan(y)))
@assert(all(!isnan(x)))
x = [ones(n) x] # append constant covariate for offset
predictor_names = ["constant",predictor_names]
p = size(x,2)

function interval(x,P)
    N = length(x)
    xsort = sort(x)
    ai = (1-P)/2
    l,u = xsort[int(floor(ai*N))], xsort[int(ceil((1-ai)*N))]
    return l,u
end

# Run sampler
nalphas = length(alphas)
use = nburn+1:n_keep
n_use = n_keep - nburn
k_posteriors = zeros(p+1,nalphas)
P_posteriors = zeros(p,nalphas)
beta_means = zeros(p,nalphas)
for (i_a,alpha) in enumerate(alphas)

    # Run sampler
    zeta = (1/n)/(1/n + 1/alpha)
    beta_r,lambda_r,keepers = VarSel.run(x,y,n_total,n_keep,zeta)
    k_r = vec(sum(beta_r.!=0, 1))
    P_nonzero = vec(mean(beta_r[:,use].!=0, 2))
    P_posteriors[:,i_a] = P_nonzero
    beta_means[:,i_a] = vec(mean(beta_r[:,use],2))

    # Compute posterior on k
    edges,counts = hist(k_r[use],-0.5:1:p+0.5)
    k_posteriors[:,i_a] = counts/n_use

    # Show traceplot of number of nonzero coefficients
    figure(1); clf(); hold(true)
    plot(k_r + 0.5*rand(length(k_r))-0.25,"k.",markersize=2)
    draw()
    savefig("cpp-traceplot-alpha=$alpha.png",dpi=150)

    # Show posterior of lambda
    figure(2); clf(); hold(true)
    plt.hist(lambda_r[use], 50)
    draw()
    savefig("cpp-lambda-alpha=$alpha.png",dpi=150)
    println("E(lambda|data) = ", mean(lambda_r[use]))

    # List of most probable nonzero coefs
    order = sortperm(P_nonzero,rev=true)

    # Show traceplots of top coefficients
    figure(3,figsize=(8,8)); clf(); hold(true)
    subplots_adjust(bottom=0.1,hspace=0.4)
    for ji = 1:8
        j = order[ji]
        subplot(8,1,ji)
        betas = vec(beta_r[j,:])
        plot(betas,"k.",markersize=2)
        xm,xM = xlim(); ym,yM = ylim()
        text((xM-xm)*0.05+xm,(yM-ym)*0.5+ym,predictor_names[j])
    end
    draw()
    savefig("cpp-traceplot-beta-alpha=$alpha.png",dpi=150)

    # Display c.d.f.s of coefficients
    figure(4,figsize=(8,8)); clf(); hold(true)
    subplots_adjust(bottom=0.1,hspace=0.4)
    for ji = 1:15
        j = order[ji]
        subplot(5,3,ji)
        betas = vec(beta_r[j,use])
        plot(sort([betas,betas]),sort([0:n_use-1,1:n_use])/n_use,"b-",markersize=0,linewidth=2)
        ylim(0,1)
        mb = max(0.01, maximum(abs(betas)))
        xlim(-mb,mb)
        xm,xM = xlim(); ym,yM = ylim()
        text((xM-xm)*0.05+xm,(yM-ym)*0.5+ym,predictor_names[j])
        l,u = interval(betas,0.95)
        plot([l,l],ylim(),"r-",markersize=0)
        plot([u,u],ylim(),"r-",markersize=0)
    end
    draw()
    savefig("cpp-beta-alpha=$alpha.png",dpi=150)

end

# Plot posteriors on k
figure(5,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
plot(0:p,VarSel.prior_k(0:p,p),"ko-",label="prior",markersize=8,linewidth=2)
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? repr(int(alpha)) : "\\infty") * "\$"
    plot(0:p,vec(k_posteriors[:,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
xticks(0:p)
xlim(0,16)
ylim(0,1)
xlabel("\$k \\,\\,\\mathrm{(\\#\\,of\\,nonzero\\,coefficients)}\$",fontsize=17)
ylabel("\$\\Pi(k|\\mathrm{data})\$",fontsize=15)
legend(numpoints=1)
draw()
savefig("cpp-k_posteriors.png",dpi=150)

# Sort coefs by overall probability of inclusion
order = sortperm(vec(mean(P_posteriors,2)),rev=true)
names_order = predictor_names[order]
for i in 1:16
    println(i,": ",names_order[i])
end

# Plot posterior probability of inclusion
figure(6,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? repr(int(alpha)) : "\\infty") * "\$"
    plot(1:p,vec(P_posteriors[order,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
d_show = 16
xticks(1:d_show)
xlim(0,d_show)
yticks(linspace(0,1,6))
ylim(0,1.1)
xlabel("\$j\\,\\,\\mathrm{(coefficient\\,index)}\$",fontsize=17)
ylabel("\$\\Pi(\\beta_j \\neq 0\\,|\\,\\mathrm{data})\$",fontsize=16)
legend(numpoints=1)
draw()
savefig("cpp-P_posteriors.png",dpi=150)

# Plot posterior means
figure(7,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? repr(int(alpha)) : "\\infty") * "\$"
    plot(1:p,vec(beta_means[order,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
d_show = 16
xticks(1:d_show)
xlim(0,d_show)
xlabel("\$j\\,\\,\\mathrm{(coefficient\\,index)}\$",fontsize=17)
ylabel("\$E(\\beta_j |\\mathrm{data})\$",fontsize=16)
legend(numpoints=1)
draw()
savefig("cpp-beta_means.png",dpi=150)

end # module

max_stars_repo_name: jwmi/CoarsenedPosterior

max_issues_repo_name: jwmi/CoarsenedPosterior

max_forks_repo_name: jwmi/CoarsenedPosterior


text:
LIGHT_SPEED_M_PER_S = 299792458.0;

struct Fmcw
    chirp_bandwidth::Any
    chirp_period::Any
    chirp_rate::Any
    sampling_period::Any
    carrier_freq::Any
end

function chirp(fmcw::Fmcw, phase_offset, n, t)
    alpha = fmcw.bandwidth / fmcw.period
    return A * cos(
        2 *
        pi *
        (
            fmcw.carrier_freq * (n * fmcw.period + t) +
            (alpha / 2) * (n * fmcw.period + t)^2 +
            phase_offset
        ),
    )
end

function if_response_approx(fmcw::Fmcw, range, velocity, amplitude, n, t)
    beat_freq = 2 * fmcw.chirp_rate * range / LIGHT_SPEED_M_PER_S
    doppler_freq = 2 * velocity * fmcw.carrier_freq / LIGHT_SPEED_M_PER_S
    phase = 2 * fmcw.carrier_freq * range / LIGHT_SPEED_M_PER_S
    return amplitude * exp(
        im * (2 * pi * ((beat_freq * t) + (doppler_freq * n * fmcw.chirp_period) + phase)),
    )
end

function if_response()
    beat_freq = 2 * fmcw.chirp_rate * range / LIGHT_SPEED_M_PER_S
    doppler_freq = 2 * velocity * fmcw.carrier_freq / LIGHT_SPEED_M_PER_S
    phase = 2 * fmcw.carrier_freq * range / LIGHT_SPEED_M_PER_S

    # TODO(omar): Implement the rest of this with all the terms
end


text:
# evs = eigenvalues
# k = orbital degeneracies
# N = num of particles
# T = temperature
function compute_occupations(evs::Vector{Float64}, k::Vector{Int}, N::Float64, T::Float64)
    @assert sum(k) >= N "$(sum(k)) states < $N electrons. increase nmax,lmax."
    mu = length(evs) == 1 ? evs[1] :
        fzero(mu -> N - dot(k,fermidirac(evs,mu,T)), evs[1], evs[end]+0.1)  # chemical potential
    f = k.*fermidirac(evs,mu,T)  # occupation numbers
    f,mu
end

function compute_occupations(evs::Vector{Float64}, k::Int, N::Float64, T::Float64)
    mu = length(evs) == 1 ? evs[1] :
        fzero(mu -> N - sum(k*fermidirac(evs,mu,T)), evs[1], evs[end]+0.1)  # chemical potential
    f = k*fermidirac(evs,mu,T)  # occupation numbers
    f,mu
end

init_mixer(mixer::AbstractMixer, grid::AbstractGrid) = return

function init_mixer(mixer::PulayMixer, grid::AbstractGrid)
    mixer.grid = grid
end

function init_scf(scf::SCF, grid::AbstractGrid)
    scf.grid = grid
    scf.converged = false
end

using Base: Callable

function dft(Z         ::Float64,
             N         ::Float64,
             f         ::Vector{Float64},
             fixed_occ ::Bool,
             nmax      ::Int,
             lmax      ::Int,
             r         ::LogarithmicGrid,
             xc        ::ldaxc,
             scf       ::SCF,
             mixer     ::AbstractMixer,
             T         ::Float64,
             Emin      ::Float64,
             Emax      ::Float64,
             Elim      ::Float64,
             logger    ::Callable)
    local E,Etot,vs,vh,evs,efs,q,mu

    stats = Stats()
    init_mixer(mixer,r)
    init_scf(scf,r)

    @stats "* total" begin

        # initialization
        if fixed_occ; N = sum(f); mu = NaN end
        rho_in = isempty(scf.rho) ? N^4*exp(-N*r/2)/64pi : scf.rho[end]         # initial density
        vext = -Z./r                                                            # external potential
        dr = 4pi*r.^2                                                           # radial integration factor

        # storage
        vx = similar(r); vc = similar(r); ex = similar(r); ec = similar(r)

        while !scf.converged

            @stats "hartree"                    vh = 4π*radpoisson(r,rho_in,0)                     # Hartree potential
            @stats "exchange-correlation"       xc(rho_in,ex,vx,ec,vc)                              # exchange-correlation
                                                vs = vext + vx + vc + vh                            # Kohn-Sham potential
            @stats "KS equations"               evs,efs,q = radschrod(r,vs,nmax,lmax,reverse=true,
                                                    Emin=Emin,Emax=Emax,Elim=Elim)
                                                l = q[:,2]                                          # l-quantum numbers
            @stats "analytic continuation"      analytic_continuation_coulomb_nuclei!(r,efs,l)

            if !fixed_occ; f,mu = compute_occupations(evs,2(2l+1),N,T) end
            @assert length(evs) >= length(f) "$(length(evs)) states < $(length(f)) required. increase nmax,lmax."

            @stats "compute density"            rho_out = (efs.^2 * f)./dr                                      # radial density
                                                E = dot(f,evs)                                                  # sum of eigenvalues
            @stats "harris energy"              Etot = E + integrate(r, dr.*rho_in.*(ex+ec - (vx+vc) - vh/2))   # Harris functional
            @stats "check convergence"          scf(Etot, rho_out)                                              # check convergence
            @stats "density mixing"             rho_in = mixer(rho_in, rho_out)                                 # mix densities

            if !is(logger,Void) logger(scf) end

        end  # scf cycle

        Ex   = integrate(r, rho_in.*ex)         # exchange energy
        Ec   = integrate(r, rho_in.*ec)         # correlation energy
        Exc  = Ex + Ec                          # exchange-correlation energy
        Eh   = integrate(r, rho_in.*vh/2)       # Hartree energy
        Epot = integrate(r, rho_in.*vext)       # potential energy (external)
        Ekin = Etot - Epot - Eh - Ex - Ec       # kinetic energy

    end  # total

    struct(Etot=Etot,Ekin=Ekin,Epot=Epot,Eh=Eh,Ex=Ex,Ec=Ec,Exc=Exc,rho=rho_in,
        vs=vs,vh=vh,evs=evs,efs=efs,q=q,N=N,f=f,mu=mu,stats=stats)
end

function dft!(at::Atom; fixed_occ=false, Z=at.Z, R=at.R, params...)
    at.Z = Z = convert(Float64,Z)
    at.R = R = convert(Point3f,R)
    merge!(at, Dict{Symbol,Any}(params))

    res = dft(
        Z,                                                      # atomic number
        get!(at, :N, Z),                                        # particle number
        fixed_occ ? at[:f] : Float64[],                         # occupations
        fixed_occ,                                              # fixed occupations
        get!(at, :nmax,   nlmax[Z][1]),                         # max n-quantum number
        get!(at, :lmax,   nlmax[Z][2]),                         # max l-quantum number
        get!(at, :r,      loggrid(atomicloggrid[Z]...)),        # radial grid
        get!(at, :xc,     ldaxcunpol()),                        # xc functional
        get!(at, :scf,    SCF()),
        get!(at, :mixer,  PulayMixer()),                        # density mixer
        fixed_occ ? NaN : get!(at, :T, 1e-6),                   # temperature
        get!(at, :Emin,   -1.),                                 # see radschrod
        get!(at, :Emax,   0.),                                  # see radschrod
        get!(at, :Elim,   1e4),                                 # see radschrod
        get!(at, :logger, Void)
    )

    merge!(at, res)
    return
end

function dft(at::Atom=Atom(); fixed_occ=false, Z=at.Z, R=at.R, params...)
    let at = deepcopy(at)
        dft!(at,fixed_occ=fixed_occ,Z=Z,R=R,params...)
        return at
    end
end

function dft(; fixed_occ=false, Z=1.0, R=(0.,0.,0.), params...)
    at = Atom(Z=Z)
    let at = deepcopy(at)
        dft!(at,fixed_occ=fixed_occ,Z=Z,R=R,params...)
        return at
    end
end

for mixer in (LinearMixer,PulayMixer)
    eval(quote
        precompile(dft, (Float64,Float64,Vector{Float64},Bool,Int,Int,LogarithmicGrid,ldaxc,SCF,$mixer,Float64,Float64,Float64,Float64,Callable))
    end)
end

precompile(dft!, (Atom,))
precompile(dft, (Atom,))

max_stars_repo_name: johncsnyder/Element


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -47.221875796417265 <= q <= 49.11122836315772)

                     add_NL_constraint(m, :(1/(1 + exp(-(0.42790162605340454 + 0.9948146339173967*1/(1 + exp(-(0.30429674417798 + 0.01326158862793747*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8506753528903843*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7699605184902198*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9481609482474158*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + 0.5690258167368487*1/(1 + exp(-(0.4003333530904172 + 0.4565361301167439*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.41218949132242644*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6719961553691296*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9454422589116072*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.632632285898699*1/(1 + exp(-(-0.4740061987673023 + -0.21503151130953757*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7641417558802628*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8796988576372375*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6466469730144713*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.0799939753544896*1/(1 + exp(-(-0.9533597279323933 + 0.019939841305541606*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.8134733780415981*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.10348468441851777*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.008542712959329624*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))))))))) + 1/(1 + exp(-(-0.31282642357163937 + 0.400776902433146*1/(1 + exp(-(0.30429674417798 + 0.01326158862793747*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8506753528903843*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7699605184902198*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9481609482474158*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.8700113611274283*1/(1 + exp(-(0.4003333530904172 + 0.4565361301167439*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.41218949132242644*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6719961553691296*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9454422589116072*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.8173535730312884*1/(1 + exp(-(-0.4740061987673023 + -0.21503151130953757*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7641417558802628*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8796988576372375*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6466469730144713*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.10383649475667989*1/(1 + exp(-(-0.9533597279323933 + 0.019939841305541606*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.8134733780415981*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.10348468441851777*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.008542712959329624*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))))))))) + 1/(1 + exp(-(0.07748802701261592 + -0.48287575350847955*1/(1 + exp(-(0.30429674417798 + 0.01326158862793747*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8506753528903843*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7699605184902198*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9481609482474158*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + 0.9032398979921337*1/(1 + exp(-(0.4003333530904172 + 0.4565361301167439*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.41218949132242644*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6719961553691296*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9454422589116072*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + 0.8154202752057489*1/(1 + exp(-(-0.4740061987673023 + -0.21503151130953757*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7641417558802628*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8796988576372375*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6466469730144713*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + 0.8885160528665201*1/(1 + exp(-(-0.9533597279323933 + 0.019939841305541606*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.8134733780415981*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.10348468441851777*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.008542712959329624*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))))))))) + 1/(1 + exp(-(0.6389491685818425 + -0.9024035564379238*1/(1 + exp(-(0.30429674417798 + 0.01326158862793747*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8506753528903843*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7699605184902198*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9481609482474158*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.8688708754259227*1/(1 + exp(-(0.4003333530904172 + 0.4565361301167439*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.41218949132242644*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6719961553691296*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.9454422589116072*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + -0.9399903753655909*1/(1 + exp(-(-0.4740061987673023 + -0.21503151130953757*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.7641417558802628*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.8796988576372375*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.6466469730144713*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2])))))))))) + 0.5656846187923494*1/(1 + exp(-(-0.9533597279323933 + 0.019939841305541606*1/(1 + exp(-(0.7020479917236249 + 0.5157446936496526*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.8162922098274223*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6507571114046486*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + 0.3903970834590602*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.8134733780415981*1/(1 + exp(-(0.7594455380381384 + 0.9088497207479063*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + 0.34867162336420376*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + -0.586942758624089*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.05285512797069236*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + -0.10348468441851777*1/(1 + exp(-(0.5333438081093944 + -0.9884706408297048*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.9213430875393112*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.6584439947423988*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.39397409418232554*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))) + 0.008542712959329624*1/(1 + exp(-(-0.8688273038537231 + -0.5762764679031651*1/(1 + exp(-(-0.990780872931377 + 0.7098282678057655*$(x[1]) + -0.9123645360112271*$(x[2])))) + -0.4115583086286181*1/(1 + exp(-(-0.26001558800187974 + -0.2578347550825866*$(x[1]) + -0.8440952761970419*$(x[2])))) + 0.07598378122288318*1/(1 + exp(-(0.6359235954642828 + 0.0835591752979501*$(x[1]) + -0.4064923799931912*$(x[2])))) + -0.5005567670062852*1/(1 + exp(-(0.919513762198326 + -0.48206842436164976*$(x[1]) + 0.2512150299638187*$(x[2]))))))))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
using Revise
using OptiMode
using LinearAlgebra, LinearMaps, IterativeSolvers, FFTW, ChainRulesCore, ChainRules, Zygote, OptiMode, BenchmarkTools
include("mpb_example.jl")
ε⁻¹ = [ε⁻¹_mpb[i,j,k][a,b] for a=1:3,b=1:3,i=1:ds.Nx,j=1:ds.Ny,k=1:ds.Nz];
ω = ω_mpb
H,k = solve_k(ω_mpb,ε⁻¹)
Nz = 1
gx = collect(fftfreq(Nx,Nx/6.0))
gy = collect(fftfreq(Ny,Ny/4.0))
gz = collect(fftfreq(Nz,Nz/1.0))
Nx,Ny,Nz = size(ε⁻¹)[end-2:end]

function foo(ω,ε⁻¹) #Δx,Δy)
	H,kz = solve_k(ω, ε⁻¹) # out = real(norm(H) * k)
	Ha = reshape(H,(2,Nx,Ny,Nz))
	# return -ω / real(dot(vec(Ha), vec(zyg_Mₖ(Ha,ε⁻¹,Zygote.hook(ā -> real(ā), kz),gx,gy,gz))))
	return -ω / real(dot(vec(Ha), vec(zyg_Mₖ(Ha,ε⁻¹,kz,gx,gy,gz))))
end

function foo2(ω,ε⁻¹,Δx,Δy,Δz)
	Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
	gx = Zygote.@ignore collect(fftfreq(Nx,Nx/Δx))
	gy = Zygote.@ignore collect(fftfreq(Ny,Ny/Δy))
	gz = Zygote.@ignore collect(fftfreq(Nz,Nz/Δz))
	H,kz = solve_k(ω, ε⁻¹) # out = real(norm(H) * k)
	Ha = reshape(H,(2,Nx,Ny,Nz))
	# return -ω / real(dot(vec(Ha), vec(zyg_Mₖ(Ha,ε⁻¹,Zygote.hook(ā -> real(ā), kz),gx,gy,gz))))
	return -ω / real(dot(vec(Ha), vec(zyg_Mₖ(Ha,ε⁻¹,kz,gx,gy,gz))))
end

########### test to fix broadcasting/allocation slowdown of ng gradient calc

function sum2(op,arr)
    return sum(op,arr)
end

function sum2adj( Δ, op, arr )
    n = length(arr)
    g = x->Δ*Zygote.gradient(op,x)[1]
    return ( nothing, map(g,arr))
end

Zygote.@adjoint function sum2(op,arr)
    return sum2(op,arr),Δ->sum2adj(Δ,op,arr)
end

Zygote.refresh()

function KpG(kz::Real,gx::Real,gy::Real,gz::Real)
	# scale = ds.kpG[i,j,k].mag
	kpg = [-gx; -gy; kz-gz]
	mag = sqrt(sum2(abs2,kpg)) # norm(kpg)
	if mag==0
		n = [0.; 1.; 0.] # SVector(0.,1.,0.)
		m = [0.; 0.; 1.] # SVector(0.,0.,1.)
	else
		if kpg[1]==0. && kpg[2]==0.    # put n in the y direction if k+G is in z
			n = [0.; 1.; 0.] #SVector(0.,1.,0.)
		else                                # otherwise, let n = z x (k+G), normalized
			ntemp = [0.; 0.; 1.] × kpg  #SVector(0.,0.,1.) × kpg
			n = ntemp / sqrt(sum2(abs2,ntemp)) # norm(ntemp) #
		end
	end
	# m = n x (k+G), normalized
	mtemp = n × kpg
	m = mtemp / sqrt(sum2(abs2,mtemp)) #norm(mtemp) # sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
	return kpg, mag, m, n
end

function KpG2(kz::Float64,gx::Float64,gy::Float64,gz::Float64)
	# scale = ds.kpG[i,j,k].mag
	kpg = SVector{3,Float64}(-gx, -gy, kz-gz)
	mag = norm(kpg)
	if mag==0
		n = SVector{3,Float64}(0.,1.,0.)
		m = SVector{3,Float64}(0.,0.,1.)
	else
		if kpg[1]==0. && kpg[2]==0.    # put n in the y direction if k+G is in z
			n = SVector{3,Float64}(0.,1.,0.)
		else                                # otherwise, let n = z x (k+G), normalized
			ntemp = SVector{3,Float64}(0.,0.,1.) × kpg
			n = ntemp / norm(ntemp) #
		end
	end
	# m = n x (k+G), normalized
	mtemp = n × kpg
	m = mtemp / norm(mtemp) # sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
	return kpg, mag, m, n
end

kpg, mag, m, n = KpG(1.5,3.5,6.5,0.0)
KpG2(1.5,3.5,6.5,0.0)
KpG(1.5,3.5,6.5,0.0)[2]
KpG2(1.5,3.5,6.5,0.0)[2]
Zygote.gradient((kz,gx,gy,gz)->KpG(kz,gx,gy,gz)[2],1.5,3.5,6.5,0.0)
Zygote.gradient((kz,gx,gy,gz)->KpG2(kz,gx,gy,gz)[2],1.5,3.5,6.5,0.0)

function ∇KpG2(a,b,c,d)
	Zygote.gradient((kz,gx,gy,gz)->KpG(kz,gx,gy,gz)[2],a,b,c,d)
end

function ∇KpG22(a,b,c,d)
	Zygote.gradient((kz::Float64,gx::Float64,gy::Float64,gz::Float64)->KpG2(kz,gx,gy,gz)[2]::Float64,a,b,c,d)
end

@btime ∇KpG2(1.5,3.5,6.5,0.0)
@btime ∇KpG22(1.5,3.5,6.5,0.0)

function new_zyg_zcross_t2c(H,kpg)
	# [-H[1] * m[2] - Hin[2] * n[2] ; Hin[1] * m[1] + Hin[2] * n[1]; 0.0 ]
	[-H[1] * kpg[3][2] - H[2] * kpg[4][2] ; H[1] * kpg[3][1] + H[2] * kpg[4][1]; 0.0 ]
end

function new_zyg_kcross_c2t(H,kpg)
	# [dot(H,kpg[4])*kpg[2];-dot(H,kpg[3])*kpg[2]]
	SVector{2,ComplexF64}(-dot(H,kpg[4])*kpg[2],dot(H,kpg[3])*kpg[2])
end

# kpgA = reshape([new_zyg_KplusG(kz,ggx,ggy,ggz) for ggx in gx, ggy in gy, ggz in gz], (1,Nx,Ny,Nz))
kpgA = [new_zyg_KplusG(kz,ggx,ggy,ggz) for ggx in gx, ggy in gy, ggz in gz]
HA = reshape(H,(2,Nx,Ny,Nz))
HSA = [SVector{2,ComplexF64}(HA[:,i,j,k]) for i=1:Nx,j=1:Ny,k=1:Nz]

d_HSA = SVector{3,ComplexF64}.([fft([dd[a] for dd in new_zyg_zcross_t2c.(HSA,kpgA)]) for a=1:3]...)
e_HSA = SVector{3,ComplexF64}.([ifft([ee[b] for ee in ε⁻¹_mpb .* d_HSA]) for b=1:3]...)
Mk_HSA = new_zyg_kcross_c2t.(e_HSA,kpgA)

# Mk_HSA = new_zyg_kcross_c2t.(SVector{3,ComplexF64}.([ifft([ee[b] for ee in ε⁻¹_mpb .* SVector{3,ComplexF64}.([fft([dd[a] for dd in new_zyg_zcross_t2c.(HSA,kpgA)]) for a=1:3]...)] ) for b=1:3]...),kpgA)
Mk_HA = reinterpret(ComplexF64,reshape(Mk_HSA,(1,Nx,Ny,Nz)))

# sum(dot.(HSA,new_zyg_kcross_c2t.(SVector{3,ComplexF64}.([ifft([ee[b] for ee in ε⁻¹_mpb .* SVector{3,ComplexF64}.([fft([dd[a] for dd in new_zyg_zcross_t2c.(HSA,kpgA)]) for a=1:3]...)] ) for b=1:3]...),kpgA) ) )

function new_zyg_Mₖ(H,ε⁻¹,kz,gx,gy,gz)
    # Hout = zeros(ComplexF64,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
	Hin = reshape(H,(2,Nx,Ny,Nz))
	H1 = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	#H2 = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = new_zyg_KplusG(kz,gx[i],gy[j],gz[k])
        H1[1,i,j,k] = -Hin[1,i,j,k] * m[2] - Hin[2,i,j,k] * n[2]
        H1[2,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1]
        H1[3,i,j,k] = 0.0
    end
	H11 = copy(H1)
    H2 = fft(H11,(2:4))
	println("$(typeof(H2))")
	H3 = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        H3[1,i,j,k] =  ε⁻¹[1,1,i,j,k]*H2[1,i,j,k] + ε⁻¹[2,1,i,j,k]*H2[2,i,j,k] + ε⁻¹[3,1,i,j,k]*H2[3,i,j,k]
        H3[2,i,j,k] =  ε⁻¹[1,2,i,j,k]*H2[1,i,j,k] + ε⁻¹[2,2,i,j,k]*H2[2,i,j,k] + ε⁻¹[3,2,i,j,k]*H2[3,i,j,k]
        H3[3,i,j,k] =  ε⁻¹[1,3,i,j,k]*H2[1,i,j,k] + ε⁻¹[2,3,i,j,k]*H2[2,i,j,k] + ε⁻¹[3,3,i,j,k]*H2[3,i,j,k]
    end
	H33 = copy(H3)
    H4 = ifft(H33,(2:4))
	Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = new_zyg_KplusG(kz,gx[i],gy[j],gz[k])
        scale = -1 / mag
        at1 = H4[1,i,j,k] * m[1] + H4[2,i,j,k] * m[2] + H4[3,i,j,k] * m[3]
        at2 = H4[1,i,j,k] * n[1] + H4[2,i,j,k] * n[2] + H4[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  at2 / mag
        Hout[2,i,j,k] =  -at1 / mag
    end
    return copy(Hout)
end

function foo3(ω,ε⁻¹,Δx,Δy,Δz)
	Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
	gx = Zygote.@ignore collect(fftfreq(Nx,Nx/Δx))
	gy = Zygote.@ignore collect(fftfreq(Ny,Ny/Δy))
	gz = Zygote.@ignore collect(fftfreq(Nz,Nz/Δz))
	H,kz = solve_k(ω, ε⁻¹) # out = real(norm(H) * k)
	Ha = reshape(H,(2,Nx,Ny,Nz))
	# return -ω / real(dot(vec(Ha), vec(zyg_Mₖ(Ha,ε⁻¹,Zygote.hook(ā -> real(ā), kz),gx,gy,gz))))
	return -ω / real(dot(vec(Ha), vec(new_zyg_Mₖ(H,ε⁻¹,kz,gx,gy,gz))))
end

function fooRD(ωA,ε⁻¹) #Δx,Δy)
	ω = ωA[1]
	H,kz = solve_k(ω, ε⁻¹) # out = real(norm(H) * k)
	Ha = reshape(H,(2,Nx,Ny,Nz))
	return -ω / real(dot(vec(Ha), vec(zyg_Mₖ(Ha,ε⁻¹,kz,gx,gy,gz))))
end

foo(ω_mpb, ε⁻¹)
foo2(ω_mpb, ε⁻¹,6.0,4.0,1.0)
foo3(ω_mpb, ε⁻¹,6.0,4.0,1.0)

Zygote.gradient(foo,ω_mpb,ε⁻¹)
Zygote.gradient(foo2,ω_mpb,ε⁻¹,6.0,4.0,1.0)
Zygote.gradient(foo3,ω_mpb,ε⁻¹,6.0,4.0,1.0)
Zygote.gradient(foo2,ω_mpb,ε⁻¹,Δx=6.0,Δy=4.0,Δz=1.0)

fooRD([ω_mpb, ] , ε⁻¹)
Zygote.refresh()
##
ng_primal, foo_pb = Zygote.pullback(foo,ω,ε⁻¹)

H̄ = rand(ComplexF64,size(ds.H⃗))[:,1]
k̄ = rand(Float64)
ΔΩ = (H̄, k̄)

foo_pb(H̄, k̄)  #(ΔΩ)

solve_k_pullback(H̄, k̄, ω_mpb, ε⁻¹_mpb;ds,neigs=1,eigind=1,maxiter=3000,tol=1e-8)

##
function zyg_ng(H,eps,kz,gx,gy,gz)
	return real(dot(vec(H), vec(zyg_Mₖ(H,eps,kz,gx,gy,gz))))
end

zyg_ng(reshape(H,(2,Nx,Ny,Nz)),ε⁻¹,k,fftfreq(Nx,Nx/6.0),fftfreq(Ny,Ny/6.0),fftfreq(Nz,Nz/1.0))

using ReverseDiff
const foo_tape = ReverseDiff.GradientTape(fooRD,([ω_mpb, ],ε⁻¹))

ChainRules.refresh_rules()
Zygote.refresh()

##

epsR = Array{Real,5}(eps);
gx = [g[1] for g in ds.gx]; gy = [g[2] for g in ds.gy]; gz = [g[3] for g in ds.gz];
Hin = copy(reshape(H,(2,ds.Nx,ds.Ny,ds.Nz)));
kz = Real(1.5)
kzA = [kz,]
kzC = ComplexF64(kz)
kzCA = [ComplexF64(kz),]
Nx,Ny,Nz = size(ε⁻¹_mpb)[end-2:end]
g = MaxwellGrid(6.0,4.0,1.0,Nx,Ny,Nz)
k_guess=ω_mpb*sqrt(1/minimum([minimum(eps[a,a,:,:,:]) for a=1:3]))
ds = MaxwellData(k_guess,g)
ε⁻¹_SHM3 = [SHM3(eps[:,:,i,j,kk]) for i=1:Nx,j=1:Ny,kk=1:Nz]
##
k_guess
##

using Roots
kz = Roots.find_zero(k -> _solve_Δω²(k,ω_mpb,ε⁻¹_SHM3,ds;neigs=1,eigind=1,maxiter=3000,tol=1e-7), ds.k[3], Roots.Newton())

_solve_Δω²(k,ω_mpb,ε⁻¹_SHM3,ds;neigs=1,eigind=1,maxiter=3000,tol=1e-7)

##
H,k = solve_k(ω_mpb,eps)

# @btime H,k = solve_k(0.66,$eps) # 998.256 ms (43540 allocations: 240.61 MiB)
# @btime H,k = solve_k(0.62,ε⁻¹_mpb,ds)    #  1.150 ms (222 allocations: 5.17 MiB)
@btime H,k = solve_k(0.62,ε⁻¹_mpb,ds)
@profiler H,k = solve_k(0.66,eps)
@profiler H,k = solve_k(0.64,ε⁻¹_mpb,ds)

# Zygote.@profile solve_k(0.66,eps)
# Zygote.@profile solve_k(0.64,ε⁻¹_mpb,ds)

##
@btime H,k = solve_k(0.66,$eps)
#

##

# calculate ng

# function foo(H::Array{Complex{Float64},2},k::Float64,ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ω::Float64)::Float64
# 	kpg = kpG(k,ds.grid::MaxwellGrid)::Array{KVec,3}
#     ng = ω /  real( ( H[:,1]' * Mₖ(H[:,1], ε⁻¹,kpg) )[1])
# end

# @profiler foo(H,k,ε⁻¹_mpb,ω_mpb)
# @btime foo($H,$k,$ε⁻¹_mpb,$ω_mpb)
# # calculate
# @benchmark Zygote.gradient(foo,$H,$k,$ε⁻¹_mpb,$ω_mpb)

# Zygote.gradient(foo,H,k,ε⁻¹_mpb,ω_mpb)

# @profiler Zygote.gradient(foo,H,k,ε⁻¹_mpb,ω_mpb)

# @profiler H̄,k̄,ε̄⁻¹,ω̄ = Zygote.gradient(foo,H,k,ε⁻¹_mpb,ω_mpb)

# Zygote.@profile foo(H,k,ε⁻¹_mpb,ω_mpb)

##

H,k = solve_k(0.62,ε⁻¹_mpb;ds)
# @btime H,k = solve_k(0.62,ε⁻¹_mpb,ds)    #  1.150 ms (222 allocations: 5.17 MiB)
nᵧ, nᵧ_back = Zygote.pullback(foo,H,ω_mpb,ε⁻¹_mpb)
H̄,ω̄,ε̄⁻¹ = nᵧ_back(1)

## calculate ng from initial modesolve inputs

# function goo(H,ω,ε⁻¹)
# 	H,k = solve_k(ω,ε⁻¹;ds)
#     ng = ω /  real( ( H[:,1]' * Mₖ(H[:,1], ε⁻¹,ds) )[1])
# end

function Mop(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    d = zcross_t2c(Hin,kpG);
    # e = ifft(ε⁻¹_dot(d,ε⁻¹),(2:4));
    kcross_c2t(d,kpG)
end

# function goo(H::Array{Complex{Float64},2},k::Float64,ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ω::Float64)::Float64
function goo(H::Array{Complex{Float64},2},k::Float64)::Float64
	# H,k = solve_k(ω,ε⁻¹;ds)
	kpg = kpG(k,ds.grid::MaxwellGrid)::Array{KVec,3}
	# ng = ω /  real( ( H[:,1]' * Mop(H[:,1], kpg) )[1])
	out = ( H[:,1]' * Mop(H[:,1], kpg) )[1]
end

@profiler goo(H,k,ε⁻¹_mpb,ω_mpb)

## Define differentiable versions of the functions used to calculate ng

# Non-mutating, all AbstractArray Operators

function KplusG(kz,gx,gy,gz)
	# scale = ds.kpG[i,j,k].mag
	kpg = [-gx; -gy; kz-gz]
	mag = sqrt(sum(abs2.(kpg)))
	if mag==0
		n = [0.; 1.; 0.] # SVector(0.,1.,0.)
		m = [0.; 0.; 1.] # SVector(0.,0.,1.)
	else
		if kpg[1]==0. && kpg[2]==0.    # put n in the y direction if k+G is in z
			n = [0.; 1.; 0.] #SVector(0.,1.,0.)
		else                                # otherwise, let n = z x (k+G), normalized
			temp = [0.; 0.; 1.] × kpg  #SVector(0.,0.,1.) × kpg
			n = temp / sqrt(sum(abs2.(temp)))
		end
	end
	# m = n x (k+G), normalized
	mtemp = n × kpg
	m = mtemp / sqrt(sum(abs2.(mtemp))) #sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
	return kpg, mag, m, n
end

function new_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
		Hout[1,i,j,k] = ( Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1] ) * mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * m[2] + Hin[2,i,j,k] * n[2] ) * mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * m[3] + Hin[2,i,j,k] * n[3] ) * mag
    end
    return Hout
end

function new_c2t(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,2,Nx,Ny,Nz) # Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
    end
    return Hout
end

function new_zcross_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = -Hin[1,i,j,k] * m[2] - Hin[2,i,j,k] * n[2]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1]
        Hout[3,i,j,k] = 0.0
    end
    return Hout
end

function new_kcross_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) * -mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) * -mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) * -mag
    end
    return Hout
end

function new_kcross_c2t(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,2,Nx,Ny,Nz) # Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  -at2 * mag
        Hout[2,i,j,k] =  at1 * mag
    end
    return Hout
end

function new_kcrossinv_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) / mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) / mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) / mag
    end
    return Hout
end

function new_kcrossinv_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,2,Nx,Ny,Nz) # Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        scale = -1 / mag
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  at2 / mag
        Hout[2,i,j,k] =  -at1 / mag
    end
    return Hout
end

function new_ε⁻¹_dot(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        Hout[1,i,j,k] =  ε⁻¹[1,1,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,1,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,1,i,j,k]*Hin[3,i,j,k]
        Hout[2,i,j,k] =  ε⁻¹[1,2,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,2,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,2,i,j,k]*Hin[3,i,j,k]
        Hout[3,i,j,k] =  ε⁻¹[1,3,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,3,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,3,i,j,k]*Hin[3,i,j,k]
        # Hout[1,i,j,k] =  ε⁻¹[i,j,k][1,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][1,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][1,3]*Hin[3,i,j,k]
        # Hout[2,i,j,k] =  ε⁻¹[i,j,k][2,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][2,3]*Hin[3,i,j,k]
        # Hout[3,i,j,k] =  ε⁻¹[i,j,k][3,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][3,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,3]*Hin[3,i,j,k]
    end
    return Hout
end

function new_ε_dot_approx(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        ε_ave = 3 / tr(ε⁻¹[:,:,i,j,k])
        Hout[1,i,j,k] =  ε_ave * Hin[1,i,j,k]
        Hout[2,i,j,k] =  ε_ave * Hin[2,i,j,k]
        Hout[3,i,j,k] =  ε_ave * Hin[3,i,j,k]
    end
    return Hout
end

function new_M(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(new_kcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(new_ε⁻¹_dot(d,ε⁻¹),(2:4)); # (-1/(π)) .*
    new_kcross_c2t(e,kz,gx,gy,gz)
end

function new_Mₖ(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(new_zcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(new_ε⁻¹_dot(d,ε⁻¹),(2:4)); # (-1/(π)) .*
    new_kcross_c2t(e,kz,gx,gy,gz)
end

# Zygote versions

function sum2(op,arr)
    return sum(op,arr)
end

function sum2adj( Δ, op, arr )
    n = length(arr)
    g = x->Δ*Zygote.gradient(op,x)[1]
    return ( nothing, map(g,arr))
end

Zygote.@adjoint function sum2(op,arr)
    return sum2(op,arr),Δ->sum2adj(Δ,op,arr)
end
Zygote.refresh()

function zyg_KplusG(kz,gx,gy,gz)
	# scale = ds.kpG[i,j,k].mag
	kpg = [-gx; -gy; kz-gz]
	mag = sqrt(sum2(abs2,kpg))
	if mag==0
		n = [0.; 1.; 0.] # SVector(0.,1.,0.)
		m = [0.; 0.; 1.] # SVector(0.,0.,1.)
	else
		if kpg[1]==0. && kpg[2]==0.    # put n in the y direction if k+G is in z
			n = [0.; 1.; 0.] #SVector(0.,1.,0.)
		else                                # otherwise, let n = z x (k+G), normalized
			ntemp = [0.; 0.; 1.] × kpg  #SVector(0.,0.,1.) × kpg
			n = ntemp / sqrt(sum2(abs2,ntemp))
		end
	end
	# m = n x (k+G), normalized
	mtemp = n × kpg
	m = mtemp / sqrt(sum2(abs2,mtemp)) #sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
	return kpg, mag, m, n
end

function zyg_t2c(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(3,size(Hin)[2:end]...))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
		Hout[1,i,j,k] = ( Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1] ) * mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * m[2] + Hin[2,i,j,k] * n[2] ) * mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * m[3] + Hin[2,i,j,k] * n[3] ) * mag
    end
    return copy(Hout)
end

function zyg_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
    end
    return copy(Hout)
end

function zyg_zcross_t2c(Hin,kz,gx,gy,gz)
    # Hout = zeros(ComplexF64,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = -Hin[1,i,j,k] * m[2] - Hin[2,i,j,k] * n[2]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1]
        Hout[3,i,j,k] = 0.0
    end
    return copy(Hout)
end

function zyg_kcross_t2c(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) * -mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) * -mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) * -mag
    end
    return copy(Hout)
end

# Hout = [ (kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k]); Hi[]) ]

function zyg_kcross_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  -at2 * mag
        Hout[2,i,j,k] =  at1 * mag
    end
    return copy(Hout)
end

function zyg_kcrossinv_t2c(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) / mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) / mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) / mag
    end
    return copy(Hout)
end

function zyg_kcrossinv_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = zyg_KplusG(kz,gx[i],gy[j],gz[k])
        scale = -1 / mag
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  at2 / mag
        Hout[2,i,j,k] =  -at1 / mag
    end
    return copy(Hout)
end

function zyg_ε⁻¹_dot(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        Hout[1,i,j,k] =  ε⁻¹[1,1,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,1,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,1,i,j,k]*Hin[3,i,j,k]
        Hout[2,i,j,k] =  ε⁻¹[1,2,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,2,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,2,i,j,k]*Hin[3,i,j,k]
        Hout[3,i,j,k] =  ε⁻¹[1,3,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,3,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,3,i,j,k]*Hin[3,i,j,k]
        # Hout[1,i,j,k] =  ε⁻¹[i,j,k][1,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][1,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][1,3]*Hin[3,i,j,k]
        # Hout[2,i,j,k] =  ε⁻¹[i,j,k][2,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][2,3]*Hin[3,i,j,k]
        # Hout[3,i,j,k] =  ε⁻¹[i,j,k][3,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][3,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,3]*Hin[3,i,j,k]
    end
    return copy(Hout)
end

function zyg_ε_dot_approx(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        ε_ave = 3 / tr(ε⁻¹[:,:,i,j,k])
        Hout[1,i,j,k] =  ε_ave * Hin[1,i,j,k]
        Hout[2,i,j,k] =  ε_ave * Hin[2,i,j,k]
        Hout[3,i,j,k] =  ε_ave * Hin[3,i,j,k]
    end
    return copy(Hout)
end

function zyg_M(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(zyg_kcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(zyg_ε⁻¹_dot(d,ε⁻¹),(2:4)); # (-1/(π)) .*
    zyg_kcross_c2t(e,kz,gx,gy,gz)
end

# function zyg_M(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData)::Array{ComplexF64,1}
#     HinA = reshape(Hin,(2,ds.Nx,ds.Ny,ds.Nz))
#     HoutA = M(HinA,ε⁻¹,ds)
#     return vec(HoutA)
# end

function zyg_Mₖ(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(zyg_zcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(zyg_ε⁻¹_dot(d,ε⁻¹),(2:4));
    zyg_kcross_c2t(e,kz,gx,gy,gz)
end

# function new_Mₖ(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},kpG::AbstractArray{KVec,3})::Array{ComplexF64,1}
#     Nx,Ny,Nz = size(ε⁻¹)
#     HinA = reshape(Hin,(2,Nx,Ny,Nz))
#     HoutA = Mₖ(HinA,ε⁻¹,kpG)
#     return -vec(HoutA)
# end

##

function ng_fm(kz::Real)::Real
    # return real(dot(vec(Hin), vec(new_Mₖ(Hin,eps,kz,gx,gy,gz))))
    return real(dot(vec(Hin), vec(new_Mₖ(Hin,eps,1.55,gx,gy,gz)))) * kz
end

function zyg_ng(H,eps,kz)
	return real(dot(vec(H), vec(zyg_Mₖ(H,eps,kz,gx,gy,gz))))
end

function zyg_ng2(kz)
	return real(dot(vec(Hin), vec(zyg_Mₖ(Hin,eps,kz,gx,gy,gz))))
end

##
#using ForwardDiff
ForwardDiff.derivative(ng_fm, kz)

# @benchmark fn(1.55)
##
# f2_inputs = Hin, epsC, kzCA
# f2_tape = ReverseDiff.GradientTape(f2, f2_inputs)
# f2_cfg = ReverseDiff.GradientConfig(f2, epsC)

#ReverseDiff.gradient(f2, (Hin, epsC, kzCA))
#@benchmark ReverseDiff.gradient(e -> f2(Hin,e,kzA),eps)
Zygote.gradient(zyg_ng,Hin,eps,1.55)

##

# function solve_k_pullback(H̄, k̄, ω, ε⁻¹) #;neigs,eigind,maxiter,tol) # ω̄ ₖ)
# 	H, kz = solve_k(ω, ε⁻¹) #;neigs,eigind,maxiter,tol) #Ω
#
# 	# hacky handling of non-differentiated parameters for now
# 	eigind = 1
# 	Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
# 	gx = fftfreq(Nx,Nx/6.0)
# 	gy = fftfreq(Ny,Ny/4.0)
# 	gz = fftfreq(Nz,Nz/1.0)
# 	## end hacky parameter handling
#
# 	P = LinearMap(x -> H[:,eigind] * dot(H[:,eigind],x),length(H[:,eigind]),ishermitian=true)
# 	A = M̂(ε⁻¹,kz,gx,gy,gz) - ω^2 * I
# 	b = ( I  -  P ) * H̄[:,eigind]
# 	λ⃗₀ = IterativeSolvers.bicgstabl(A,b,3)
# 	ωₖ = real( ( H[:,eigind]' * Mₖ(H[:,eigind], ε⁻¹,kz,gx,gy,gz) )[1]) / ω # ds.ω²ₖ / ( 2 * ω )
# 	# Hₖ =  ( I  -  P ) * * ( Mₖ(H[:,eigind], ε⁻¹,ds) / ω )
# 	ω̄  =  ωₖ * k̄
# 	λ⃗₀ -= P*λ⃗₀ - ω̄  * H
# 	Ha = reshape(H,(2,ds.Nx,ds.Ny,ds.Nz))
# 	Ha_F =  ds.𝓕 * kcross_t2c(Ha,ds.kpG)
# 	λ₀ = reshape(λ⃗₀,(2,ds.Nx,ds.Ny,ds.Nz))
# 	λ₀_F  = ds.𝓕 * kcross_t2c(λ₀,ds.kpG)
# 	# ε̄ ⁻¹ = ( ds.𝓕 * kcross_t2c(λ₀,ds) ) .* ( ds.𝓕 * kcross_t2c(Ha,ds) )
# 	ε⁻¹_bar = [ Diagonal( real.(λ₀_F[:,i,j,kk] .* Ha_F[:,i,j,kk]) ) for i=1:ds.Nx,j=1:ds.Ny,kk=1:ds.Nz]
# 	return ω̄ , ε⁻¹_bar
# end

# Zygote.@adjoint solve_k(ω,ε⁻¹) = solve_k(ω,ε⁻¹), (H̄, k̄ ) -> solve_k_pullback(H̄, k̄, ω,ε⁻¹)

Zygote.@adjoint function solve_k(ω,ε⁻¹)
	H, kz = solve_k(ω,ε⁻¹)
	(H, kz), (H̄, k̄ ) -> begin
		# hacky handling of non-differentiated parameters for now
		eigind = 1
		Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
		gx = fftfreq(Nx,Nx/6.0)
		gy = fftfreq(Ny,Ny/4.0)
		gz = fftfreq(Nz,Nz/1.0)
		## end hacky parameter handling

		P = LinearMap(x -> H[:,eigind] * dot(H[:,eigind],x),length(H[:,eigind]),ishermitian=true)
		A = M̂(ε⁻¹,kz,gx,gy,gz) - ω^2 * I
		b = ( I  -  P ) * H̄[:,eigind]
		λ⃗₀ = IterativeSolvers.bicgstabl(A,b,3)
		ωₖ = real( ( H[:,eigind]' * Mₖ(H[:,eigind], ε⁻¹,kz,gx,gy,gz) )[1]) / ω # ds.ω²ₖ / ( 2 * ω )
		# Hₖ =  ( I  -  P ) * * ( Mₖ(H[:,eigind], ε⁻¹,ds) / ω )
		ω̄  =  ωₖ * k̄
		λ⃗₀ -= P*λ⃗₀ - ω̄  * H
		Ha = reshape(H,(2,Nx,Ny,Nz))
		Ha_F =  fft(kcross_t2c(Ha,kz,gx,gy,gz),(2:4))
		λ₀ = reshape(λ⃗₀,(2,Nx,Ny,Nz))
		λ₀_F  = fft(kcross_t2c(λ₀,kz,gx,gy,gz),(2:4))
		# ε̄ ⁻¹ = ( 𝓕 * kcross_t2c(λ₀,ds) ) .* ( 𝓕 * kcross_t2c(Ha,ds) )
		ε⁻¹_bar = [ Diagonal( real.(λ₀_F[:,i,j,kk] .* Ha_F[:,i,j,kk]) ) for i=1:Nx,j=1:Ny,kk=1:Nz]
		return ω̄ , ε⁻¹_bar
	end
end

# Zygote.@adjoint function vcat(xs::Union{Number, AbstractVector}...)
#     vcat(xs...), dy -> begin
#         d = 0
#         map(xs) do x
#             x isa Number && return dy[d+=1]
#             l = length(x)
#             dx = dy[d+1:d+l]
#             d += l
#             return dx
#         end
#     end
# end

Zygote.refresh()

##
LinearAlgebra.ldiv!(c,A::LinearMaps.LinearMap,b) = mul!(c,A',b)

H, kz = solve_k(ω_mpb, ε⁻¹_mpb;ds)
H = H[:,1]
P = LinearMap(x -> H * dot(H,x),length(H),ishermitian=true)
# ⟂ = I  -  P #LinearMap(x -> H * dot(H,x),length(H),ishermitian=true)
A = M̂(ε⁻¹_mpb;ds) - ω_mpb^2 * I
println(typeof(A))
b = ( I  -  P ) * H̄
λ⃗₀ = IterativeSolvers.bicgstabl(A,b,2)
# Hₖ =  ( I  -  P ) * * ( Mₖ(H[:,eigind], ε⁻¹,ds) / ω )
eigind = 1
ωₖ = real( ( H[:,eigind]' * Mₖ(H[:,eigind], ε⁻¹_mpb,ds) )[1]) / ω_mpb # ω²ₖ / ( 2 * ω )
ω̄  =  ωₖ * k̄

λ⃗₀ -= P*λ⃗₀ - ω̄  * H
Ha = reshape(H,(2,ds.Nx,ds.Ny,ds.Nz))
λ₀ = reshape(λ⃗₀,(2,ds.Nx,ds.Ny,ds.Nz))
# ε̄ ⁻¹ = ( ds.𝓕 * kcross_t2c(λ₀,ds) ) .* ( ds.𝓕 * kcross_t2c(Ha,ds) )
ε⁻¹_bar = [ Diagonal( real.( ( ds.𝓕 * kcross_t2c(λ₀,ds.kpG) )[:,i,j,kk] .* ( ds.𝓕 * kcross_t2c(Ha,ds.kpG) )[:,i,j,kk] ) ) for i=1:ds.Nx,j=1:ds.Ny,kk=1:ds.Nz]

##
H̄ = rand(ComplexF64,size(ds.H⃗))[:,1]
k̄ = rand(Float64)
solve_k_pullback(H̄, k̄, ω_mpb, ε⁻¹_mpb;ds,neigs=1,eigind=1,maxiter=3000,tol=1e-8)

##

@code_typed solve_k(ω_mpb,ε⁻¹_mpb;ds)
##

Zygote.gradient(solve_k,ω_mpb,eps)
# Zygote.gradient(solve_k,ω_mpb,ε⁻¹_mpb,ds)
Zygote.pullback(solve_k,ω_mpb,eps)

##
function solve_k2(ω::Float64,ε⁻¹::Array{SHM3,3};neigs=1,eigind=1,maxiter=3000,tol=1e-8)
    kz = Roots.find_zero(k -> _solve_Δω²(k,ω,ε⁻¹,ds;neigs,eigind,maxiter,tol), ds.k[3], Roots.Newton())
    # ds.ω = √ds.ω²
    # ds.ωₖ = ds.ω²ₖ / ( 2 * ds.ω )
    return ds.H⃗, kz #, 1/ds.ωₖ
end

Juno.@enter Zygote.gradient((om,ep)->sum(abs2.(solve_k(om,ep)[1])),ω_mpb,eps)

Zygote.gradient(ω_mpb,ε⁻¹) do ω, εinv
	H,k = solve_k(ω, εinv)
	real( sum(abs2.(H)) + k )
end

##
# zyg_ng2(1.55)
# Zygote.gradient(zyg_ng2,1.55)
Zygote.refresh()
@benchmark Zygote.gradient(zyg_ng2,1.55)
##
# Zygote.gradient(H->zyg_ng(H,1.55),Hin)
@benchmark Zygote.gradient(e->zyg_ng($Hin,1.55,e),$eps)

##

Zygote.refresh()

@benchmark Zygote.gradient(zyg_ng,$Hin,1.55)

# @benchmark Zygote.gradient(f2,$Hin,$eps,1.55)

# BenchmarkTools.Trial:
#   memory estimate:  151.14 GiB
#   allocs estimate:  8524960
#   --------------
#   minimum time:     31.358 s (15.62% GC)
#   median time:      31.358 s (15.62% GC)
#   mean time:        31.358 s (15.62% GC)
#   maximum time:     31.358 s (15.62% GC)
#   --------------
#   samples:          1
#   evals/sample:     1

gx = collect(fftfreq(Nx,Nx/6.0))
gy = collect(fftfreq(Ny,Ny/4.0))
gz = collect(fftfreq(Nz,Nz/1.0))
## end hacky parameter handling

eigind=1
kz = k
P = LinearMap(x -> H[:,eigind] * dot(H[:,eigind],x),length(H[:,eigind]),ishermitian=true)
A = M̂(ε⁻¹,kz,gx,gy,gz) - ω^2 * I
b = ( I  -  P ) * H̄[:,eigind]
λ⃗₀ = IterativeSolvers.bicgstabl(A,b,3)
ωₖ = real( ( H[:,eigind]' * Mₖ(H[:,eigind], ε⁻¹,kz,gx,gy,gz) )[1]) / ω # ds.ω²ₖ / ( 2 * ω )
# Hₖ =  ( I  -  P ) * * ( Mₖ(H[:,eigind], ε⁻¹,ds) / ω )
ω̄  =  ωₖ * real(k̄)
λ⃗₀ -= P*λ⃗₀ - ω̄  * H
Ha = reshape(H,(2,Nx,Ny,Nz))
Ha_F =  fft(kcross_t2c(Ha,kz,gx,gy,gz),(2:4))
λ₀ = reshape(λ⃗₀,(2,Nx,Ny,Nz))
λ₀_F  = fft(kcross_t2c(λ₀,kz,gx,gy,gz),(2:4))
# ε̄ ⁻¹ = ( 𝓕 * kcross_t2c(λ₀,ds) ) .* ( 𝓕 * kcross_t2c(Ha,ds) )
ε⁻¹_bar = [ Diagonal( real.(λ₀_F[:,i,j,kk] .* Ha_F[:,i,j,kk]) ) for i=1:Nx,j=1:Ny,kk=1:Nz]

ε⁻¹_bar = [ Diagonal( real.(λ₀_F[:,i,j,kk] .* Ha_F[:,i,j,kk]) )[a,b] for a=1:3,b=1:3,i=1:Nx,j=1:Ny,kk=1:Nz]