text:
module Bodies

import Base:diff,length,push!,vec

using Statistics: mean

export Body,RigidTransform,midpoints,dlength,normal,dlengthmid,centraldiff,normalmid

abstract type Body{N} end

include("bodies/bodylist.jl")

"""
    RigidTransform(x::Tuple{Float64,Float64},α::Float64)

Construct a rigid-body transform operator, with rotation by angle `α` and
translation specified by `x`. The translation coordinates are specified in the
target coordinate system.

The resulting transform can be used as an operator on pairs of coordinate vectors,
`x` and `y`, or on bodies. For transformation of bodies, it only overwrites the
`x` and `y` fields of the body, but leaves the `x̃` and `ỹ` (body coordinates) intact.

The translation can be provided as either a tuple `(x,y)` or as a complex number.

# Constructors

- `RigidTransform((x,y),α)`
- `RigidTransform(u::Vector{Float64})`
- `RigidTransform(u::NTuple{3,Float64})`
- `RigidTransform.(u)` where `u` is a collection of vectors or tuples.

# Example

```jldoctest
julia> body = Bodies.Ellipse(0.5,0.1,100)
Elliptical body with 100 points and semi-axes (0.5,0.1)
   Current position: (0.0,0.0)
   Current angle (rad): 0.0

julia> T = RigidTransform((1.0,1.0),π/4)
Rigid-body transform
  Translation: (1.0,1.0)
  Rotation angle (rad): 0.7853981633974483

julia> T(body)
Elliptical body with 100 points and semi-axes (0.5,0.1)
   Current position: (1.0,1.0)
   Current angle (rad): 0.7853981633974483
```
"""
struct RigidTransform
   α   :: Float64
   rot :: Matrix{Float64}
   trans  :: Tuple{Float64,Float64}
end

function RigidTransform(x::Tuple{Float64,Float64},α::Float64)
    rot = [cos(α) -sin(α)
           sin(α) cos(α)]
    RigidTransform(α,rot,x)
end

RigidTransform(x::Union{Vector{Float64},NTuple{3,Float64}}) = RigidTransform((x[1],x[2]),x[3])

RigidTransform(c::ComplexF64,α::Float64) = RigidTransform((real(c),imag(c)),α)

#RigidTransform(x::Vector{Vector{Float64}}) =

function Base.show(io::IO, T::RigidTransform)
    name = "Rigid-body transform"
    println(io, name)
    println(io, "  Translation: ($(T.trans[1]),$(T.trans[2]))")
    println(io, "  Rotation angle (rad): $(T.α)")
end

"""
    vec(T::RigidTransform) -> Vector{Float64}

Returns a length-3 vector of the form [x,y,α] corresponding to the translation
and rotation specified by the given transform `T`.
"""
vec(T::RigidTransform) = [T.trans[1],T.trans[2],T.α]

"""
    vec(tl::Vector{RigidTransform}) -> NTuple{N,Vector{Float64}}

Returns a tuple of length-3 vectors of the form [x,y,α] corresponding to the translation
and rotation specified by the given by the list of transforms `tl`.
"""
function vec(tl::Union{Vector{RigidTransform},NTuple{N,RigidTransform}}) where {N}

  u = (vec(tl[1]),)
  for i = 2:length(tl)
    u = u...,vec(tl[i])
  end
  return u

end

function (T::RigidTransform)(x̃::Float64,ỹ::Float64)
    Xr = T.rot*[x̃,ỹ]
    return T.trans .+ (Xr[1],Xr[2])
end
function (T::RigidTransform)(x̃::AbstractVector{Float64},ỹ::AbstractVector{Float64})
    x = deepcopy(x̃)
    y = deepcopy(ỹ)
    for i = 1:length(x̃)
        x[i],y[i] = T(x̃[i],ỹ[i])
    end
    return x, y
end

function (T::RigidTransform)(b::Body{N}) where {N}
  b.x, b.y = T(b.x̃,b.ỹ)
  b.α = T.α
  b.cent = T.trans
  return b
end

# Evaluate some geometric details of a body
"""
    length(body::Body)

Return the number of points on the body perimeter
"""
length(::Body{N}) where {N} = N

"""
    diff(body::Body/BodyList) -> Tuple{Vector{Float64},Vector{Float64}}

Compute the x and y differences of the faces on the perimeter of body `body`, whose ends
are at the current `x` and `y` coordinates (in inertial space) of the body. Face 1
corresponds to the face between points 1 and 2, for example.

If `body` is a `BodyList`, then it computes the differences separately on each
constituent body.
"""
diff(b::Body{N}) where {N} = _diff(b.x,b.y)

function diff(bl::BodyList)
    dx = Float64[]
    dy = Float64[]
    for b in bl
        dxb, dyb = diff(b)
        append!(dx,dxb)
        append!(dy,dyb)
    end
    return dx, dy
end

function _diff(x::Vector{Float64},y::Vector{Float64})
  N = length(x)
  @assert N == length(y)

  ip1(i) = 1+mod(i,N)
  dxtmp = [x[ip1(i)] - x[i] for i = 1:N]
  dytmp = [y[ip1(i)] - y[i] for i = 1:N]

  return dxtmp, dytmp

end

"""
    midpoints(body::Body/BodyList) -> Tuple{Vector{Float64},Vector{Float64}}

Compute the x and y midpoints of the faces on the perimeter of body `body`, whose ends
are at the current `x` and `y` coordinates (in inertial space) of the body. Face 1
corresponds to the face between points 1 and 2, for example.

If `body` is a `BodyList`, then it computes the differences separately on each
constituent body.
"""
midpoints(b::Body{N}) where {N} = _midpoints(b.x,b.y)

function midpoints(bl::BodyList)
    xc = Float64[]
    yc = Float64[]
    for b in bl
        xcb, ycb = midpoints(b)
        append!(xc,xcb)
        append!(yc,ycb)
    end
    return xc, yc
end

function _midpoints(x::Vector{Float64},y::Vector{Float64})

  N = length(x)
  @assert N == length(y)

  ip1(i) = 1+mod(i,N)
  xc = 0.5*[x[ip1(i)] + x[i] for i = 1:N]
  yc = 0.5*[y[ip1(i)] + y[i] for i = 1:N]

  return xc, yc

end

"""
    centraldiff(body::Body/BodyList) -> Tuple{Vector{Float64},Vector{Float64}}

Compute the circular central differences of coordinates on body `body` (or
on each body in list `body`).
"""
function centraldiff(b::Body{N}) where {N}
  ip1(i) = 1+mod(i,N)
  xc, yc = midpoints(b)
  xc .= circshift(xc,1)
  yc .= circshift(yc,1)

  return _diff(xc,yc)
end

function centraldiff(bl::BodyList)
    dx = Float64[]
    dy = Float64[]
    for b in bl
        dxb, dyb = centraldiff(b)
        append!(dx,dxb)
        append!(dy,dyb)
    end
    return dx, dy
end

"""
    dlength(body::Body/BodyList) -> Vector{Float64}

Compute the lengths of the faces on the perimeter of body `body`, whose ends
are at the current `x` and `y` coordinates (in inertial space) of the body. Face 1
corresponds to the face between points 1 and 2, for example.
"""
function dlength(b::Union{Body,BodyList})
  dx, dy = diff(b)
  return sqrt.(dx.^2+dy.^2)
end

"""
    dlengthmid(body::Body/BodyList) -> Vector{Float64}

Compute the lengths of the faces formed between the face midpoints on the
perimeter of body `body`. The indexing of these midpoint faces is consistent
with that of the regular vertex points adjacent to both midpoints.
Midpoint face 2 corresponds to the face between midpoints 1 and 2, for example.
"""
function dlengthmid(b::Union{Body,BodyList})
  dx, dy = centraldiff(b)
  return sqrt.(dx.^2+dy.^2)
end

"""
    normal(body::Body/BodyList) -> Tuple{Vector{Float64},Vector{Float64}}

Compute the current normals (in inertial components) of the faces on the perimeter
of body `body`, whose ends are at the current `x` and `y` coordinates (in inertial space)
of the body. Face 1 corresponds to the face between points 1 and 2, for example.
"""
function normal(b::Union{Body,BodyList})
  dx, dy = diff(b)
  ds = dlength(b)
  return -dy./ds, dx./ds
end

"""
    normalmid(body::Body/BodyList) -> Tuple{Vector{Float64},Vector{Float64}}

Compute the current normals (in inertial components) of the faces formed between
midpoints on the perimeter of body `body` (or each body in list `body`).
"""
function normalmid(b::Union{Body,BodyList})
  dx, dy = centraldiff(b)
  ds = dlengthmid(b)
  return -dy./ds, dx./ds
end

include("bodies/shapes.jl")

end


text:
#!/usr/bin/env julia

const DATA_DIR   = "../12_data"
const DEG2RAD    = pi / 180.0
const INPUT_FILE = "input.txt"

# ----------------------------------------
macro assert(bool_expr)
    message = string("Assertion: <<$bool_expr>> failed")
    :($(esc(bool_expr)) || error($message))
end
# ----------------------------------------

function manhattan_dist(p)
    return abs(p[1]) + abs(p[2])
end

function part1(cmds)
    println("========== PART 1 ==========")
    state_pos = [0.0, 0.0]
    state_dir = 0.0

    for (opt, val) in cmds
        # println("Part1: opt, val=$opt, $val")
        if opt == 'N';     state_pos[2] += val
        elseif opt == 'S'; state_pos[2] -= val
        elseif opt == 'W'; state_pos[1] -= val
        elseif opt == 'E'; state_pos[1] += val
        elseif opt == 'L'; state_dir += val
        elseif opt == 'R'; state_dir -= val
        elseif opt == 'F'
            dx = val * cos(state_dir * pi / 180.0)
            dy = val * sin(state_dir * pi / 180.0)
            state_pos += [dx, dy]
        end
        # println("\t", "Part 1: Current pos=$(state_pos[1]),$(state_pos[2]), dir=$(state_dir)")
    end
    dist = round(manhattan_dist(state_pos), digits=2)
    println("Part 1: Distance from original location = $(dist)")
end

function part2(cmds)
    println()
    println("========== PART 2 ==========")
    state_pos = [0.0, 0.0]
    state_dir = [10.0, 1.0]

    for (opt, val) in cmds
        # println("Part 2: opt=$opt, val=$(val)")
        if opt == 'N';     state_dir[2] += val
        elseif opt == 'S'; state_dir[2] -= val
        elseif opt == 'W'; state_dir[1] -= val
        elseif opt == 'E'; state_dir[1] += val
        elseif opt == 'L' || opt == 'R'
            cs = cos(val * DEG2RAD)
            sn = sin(val * DEG2RAD)
            if (opt == 'R')
                sn = -1.0 * sn
            end
            new_dir = [cs -sn; sn cs] * state_dir
            state_dir = new_dir
        elseif opt == 'R'
            state_dir -= val
        elseif opt == 'F'
            dx = val * state_dir[1]
            dy = val * state_dir[2]
            state_pos += [dx, dy]
        end
        # println("\tPart 2: Current pos=$(state_pos[1]),$(state_pos[2]), dir=$(state_dir)")
    end
    dist = round(manhattan_dist(state_pos), digits=2)
    println("Part 2: Distance from original location = $(dist)")
end

# ----------------------------------------

function main()
    filename = "$DATA_DIR/$INPUT_FILE"
    lines = readlines(filename)
    cmds = map(s->(s[1], parse(Float64, s[2:end])), lines)

    part1(cmds)
    part2(cmds)
end

main()

ext: jl

max_stars_repo_name: jaycoskey/AdventOfCode

max_issues_repo_name: jaycoskey/AdventOfCode

max_forks_repo_name: jaycoskey/AdventOfCode


text:
# utils.jl

using PyCall
using SparseArrays
using LinearAlgebra

const scipy_sparse_find = pyimport("scipy.sparse")["find"]
function jlsparse(Apy::PyObject)
    IA, JA, SA = scipy_sparse_find(Apy)
    return sparse(Int[i+1 for i in IA], Int[i+1 for i in JA], SA)
end

function jlload_data(dataset_str)
    # Loads input data from gcn/data directory
py"""
import sys
import pickle as pkl
import numpy as np
import networkx as nx
import scipy.sparse as sp

def parse_index_file(filename):
    index = []
    for line in open(filename):
        index.append(int(line.strip()))
    return index

def sample_mask(idx, l):
    mask = np.zeros(l)
    mask[idx] = 1
    return np.array(mask, dtype=np.bool)

def load_data(dataset_str):
    names = ['x', 'y', 'tx', 'ty', 'allx', 'ally', 'graph']
    objects = []
    for i in range(len(names)):
        with open("data/ind.{}.{}".format(dataset_str, names[i]), 'rb') as f:
            if sys.version_info > (3, 0):
                objects.append(pkl.load(f, encoding='latin1'))
            else:
                objects.append(pkl.load(f))

    x, y, tx, ty, allx, ally, graph = tuple(objects)
    test_idx_reorder = parse_index_file("data/ind.{}.test.index".format(dataset_str))
    test_idx_range = np.sort(test_idx_reorder)

    if dataset_str == 'citeseer':
        # Fix citeseer dataset (there are some isolated nodes in the graph)
        # Find isolated nodes, add them as zero-vecs into the right position
        test_idx_range_full = range(min(test_idx_reorder), max(test_idx_reorder)+1)
        tx_extended = sp.lil_matrix((len(test_idx_range_full), x.shape[1]))
        tx_extended[test_idx_range-min(test_idx_range), :] = tx
        tx = tx_extended
        ty_extended = np.zeros((len(test_idx_range_full), y.shape[1]))
        ty_extended[test_idx_range-min(test_idx_range), :] = ty
        ty = ty_extended

    features = sp.vstack((allx, tx)).tolil()
    features[test_idx_reorder, :] = features[test_idx_range, :]
    adj = nx.adjacency_matrix(nx.from_dict_of_lists(graph))

    labels = np.vstack((ally, ty))
    labels[test_idx_reorder, :] = labels[test_idx_range, :]

    idx_test = test_idx_range.tolist()
    idx_train = range(len(y))
    idx_val = range(len(y), len(y)+500)

    train_mask = sample_mask(idx_train, labels.shape[0])
    val_mask = sample_mask(idx_val, labels.shape[0])
    test_mask = sample_mask(idx_test, labels.shape[0])

    y_train = np.zeros(labels.shape)
    y_val = np.zeros(labels.shape)
    y_test = np.zeros(labels.shape)
    y_train[train_mask, :] = labels[train_mask, :]
    y_val[val_mask, :] = labels[val_mask, :]
    y_test[test_mask, :] = labels[test_mask, :]

    return adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask"""

    adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask = py"load_data"(dataset_str)
    features = jlsparse(features)
    adj = jlsparse(adj)

    return adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask
end

function normalize_adj(adj)
    # Symmetrically normalize adjacency matrix.
    rowsum = sum(adj, dims=2)
    d_inv_sqrt = vec(rowsum.^(-.5))
    if any(d_inv_sqrt .== Inf)
        d_inv_sqrt[isinf.(d_inv_sqrt)] = 0.0
    end
    d_mat_inv_sqrt = spdiagm(0 => d_inv_sqrt)
    return (adj * d_mat_inv_sqrt)' * d_mat_inv_sqrt
end

function preprocess_adj(adj)
    # Preprocessing of adjacency matrix for simple GCN model and conversion to tuple representation.
    adj_normalized = normalize_adj(adj) + sparse(1*I, size(adj))
    return adj_normalized #sparse_to_tuple
end

function masked_accuracy(preds, labels, mask)
    # accuracy with masking
    correct_pred = equal(argmax(preds, 1), argmax(labels, 1))
    accuracy_all = cast(correct_pred, Float32)
    mask = cast(mask, dtype=Float32)
    mask /= reduce_mean(mask)
    accuracy_all *= mask
    return reduce_mean(accuracy_all)
end

max_stars_repo_name: mcbaron/fluxGCN

max_issues_repo_name: mcbaron/fluxGCN

max_forks_repo_name: mcbaron/fluxGCN


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/BoundingSphere.jl/blob/master/LICENSE

@testset "2 points 2d" begin

    ball_ref = MB.SqBall{Array{Float64,1},Float64}([1.79979, -0.419288], 38.387992027461046)
    pts = StaticArrays.SArray{Tuple{2},Float64,1,2}[[0.379453, 0.65952], [1.29807, 1.06029], [1.66662, -0.225889]]

    c,r = boundingsphere(pts, WelzlPivot())
    ball = MB.SqBall(c, r^2)
    @test MB.allinside(pts, ball, rtol=1e-1)
    @test MB.allinside(pts, ball, rtol=1e-6)

    # not sure if WelzlMTF should give a good result
    c,r = boundingsphere(pts, WelzlMTF())
    ball = MB.SqBall(c, r^2)
    @test MB.allinside(pts, ball, rtol=1e-1)
    @test MB.allinside(pts, ball, rtol=1e-6)
end

@testset "2 points 1d" begin

    x = -5.
    y = x + 10*eps(Float64)
    pts = [[x], [y]]

    c,r = boundingsphere(pts, WelzlPivot())
    ball = MB.SqBall(c, r^2)
    @test MB.allinside(pts, ball, atol=100eps(Float64))

    # not sure if WelzlMTF should give a good result
    @test_broken boundingsphere(pts, WelzlMTF())
    # ball = MB.SqBall(c, r^2)
    # @test_broken MB.allinside(pts, ball, rtol=1e-1)
    # @test_broken MB.allinside(pts, ball, rtol=1e-6)
end

@testset "3 points 3d" begin

    ball_ref = MB.SqBall{Array{Float64,1},Float64}([1.05415, -2.52996, -0.584979], 20.953846606252085)
    pts = StaticArrays.SArray{Tuple{3},Float64,1,3}[[1.18024, 0.0853978, -3.01374], [0.20966, -3.69213, 3.76129], [4.51103, -0.881877, 1.92254]]

    c,r = boundingsphere(pts, WelzlPivot())
    ball = MB.SqBall(c, r^2)
    @test MB.allinside(pts, ball, rtol=1e-1)
    @test MB.allinside(pts, ball, rtol=1e-6)

    # not sure if WelzlMTF should give a good result
    c,r = boundingsphere(pts, WelzlMTF())
    ball = MB.SqBall(c, r^2)
    @test MB.allinside(pts, ball, rtol=1e-1)
    @test MB.allinside(pts, ball, rtol=1e-6)
end


text:
export CoreTrackerStatus,
    is_success,
    is_terminated,
    is_tracking,
    is_invalid_startvalue,
    CoreTrackerResult,
    solution,
    CoreTrackerOptions,
    CoreTracker,
    coretracker,
    coretracker_startsolutions,
    affine_tracking,
    track,
    track!,
    init!,
    iterator,
    current_x,
    current_t,
    current_Δt,
    steps,
    iters,
    status,
    accuracy,
    max_corrector_iters,
    max_step_size,
    set_accuracy!,
    set_max_corrector_iters!,
    set_max_step_size!,
    options,
       # deprecated
    digits_lost,
    setup!,
    refinement_accuracy,
    max_refinement_iters,
    set_max_refinement_iters!,
    set_refinement_accuracy!,
    start_parameters!,
    target_parameters!

const coretracker_supported_keywords = [
    :accuracy,
    :auto_scaling,
    :initial_step_size,
    :min_step_size,
    :max_cond,
    :max_corrector_iters,
    :max_step_size,
    :max_steps,
    :norm,
    :patch,
    :precision,
    :predictor,
    :simple_step_size_alg,
    :terminate_ill_conditioned,
    :log_homotopy,
    :log_transform,
    :logarithmic_time_scale,
    :from_infinity,
    # deprecated,
    :max_lost_digits,
    :max_refinement_iters,
    :refinement_accuracy,
]

####################
# CoreTrackerState #
####################

module CoreTrackerStatus

@doc """
    CoreTrackerStatus.states

The possible states a `CoreTracker` can have are

* `CoreTrackerStatus.success`: Indicates a successfull completed tracking
* `CoreTrackerStatus.tracking`: The tracking is still in progress
* `CoreTrackerStatus.terminated_maximal_iterations`: Tracking terminated since maximal iterations reached.
* `CoreTrackerStatus.terminated_ill_conditioned`: Tracking terminated since the path was too ill-conditioned.
* `CoreTrackerStatus.terminated_invalid_startvalue`: Tracking terminated since the provided start value was invalid.
* `CoreTrackerStatus.terminated_step_size_too_small`
"""
@enum states begin
    success
    tracking
    terminated_maximal_iterations
    terminated_accuracy_limit
    terminated_ill_conditioned
    terminated_invalid_startvalue
    terminated_step_size_too_small
        # no more used
    terminated_singularity
end
end

"""
    is_success(S::CoreTrackerStatus.states)

Returns `true` if `S` indicates a success in the path tracking.
"""
is_success(S) = S == CoreTrackerStatus.success

"""
    is_terminated(S::CoreTrackerStatus.states)

Returns `true` if `S` indicates that the path tracking got terminated. This is not `true`
if `is_success(S)` is `true`.
"""
is_terminated(S::CoreTrackerStatus.states) =
    S ≠ CoreTrackerStatus.tracking && S ≠ CoreTrackerStatus.success

"""
    is_invalid_startvalue(S::CoreTrackerStatus.states)

Returns `true` if `S` indicates that the path tracking got terminated since the start
value was not a zero.
"""
is_invalid_startvalue(S::CoreTrackerStatus.states) =
    S == CoreTrackerStatus.terminated_invalid_startvalue

"""
    is_tracking(S::CoreTrackerStatus.states)

Returns `true` if `S` indicates that the path tracking is not yet finished.
"""
is_tracking(S::CoreTrackerStatus.states) = S == CoreTrackerStatus.tracking

####################
# CoreTrackerResult #
####################

"""
     CoreTrackerResult{V<:AbstractVector}

The result of [`track(::CoreTracker, x₁, t₁, t₀)`](@ref). You can use
[`is_success`](@ref) to check whether the tracking was successfull and [`solution`](@ref)
to obtain the solution.

# Fields

* `returncode` The [`CoreTrackerStatus.states`](@ref) enum.
* `x::V` The solution at `t`.
* `t::ComplexF64` The time `t` when the path tracker stopped.
* `accuracy::Float64`: The estimated accuracy of `x`.
* `accepted_steps::Int`: The number of accepted steps during the tracking.
* `rejected_steps::Int`: The number of rejected_steps steps during the tracking.
"""
struct CoreTrackerResult{V<:AbstractVector}
    returncode::CoreTrackerStatus.states
    x::V
    t::ComplexF64
    accuracy::Float64
    accepted_steps::Int
    rejected_steps::Int
end

function CoreTrackerResult(tracker)
    state = tracker.state
    CoreTrackerResult(
        state.status,
        copy(state.x),
        state.segment[state.s],
        state.accuracy,
        state.accepted_steps,
        state.rejected_steps,
    )
end

"""
    is_success(result::CoreTrackerResult)

Returns `true` if the path tracking was successfull.
"""
is_success(result::CoreTrackerResult) = is_success(result.returncode)

"""
    solution(result::CoreTrackerResult)

Returns the solutions obtained by the `CoreTracker`.
"""
solution(result::CoreTrackerResult) = result.x

steps(result::CoreTrackerResult) = result.accepted_steps + result.rejected_steps

Base.show(io::IO, result::CoreTrackerResult) = print_fieldnames(io, result)
Base.show(io::IO, ::MIME"application/prs.juno.inline", result::CoreTrackerResult) = result

######################
# CoreTrackerOptions #
######################
"""
    PrecisionOption

Controls the used precision for computing the residual in Newton's method.
See [[Tisseur01]](https://epubs.siam.org/doi/abs/10.1137/S0895479899359837) for the analysis behind this approach.

## Values
* `PRECISION_FIXED_64`: Only use default 64 bit machine precision
* `PRECISION_FIXED_128`: Always use emulated 128 bit floating point numbers. These are provided by the [DoubleFloats.jl](https://github.com/JuliaMath/DoubleFloats.jl) package.
* `PRECISION_ADAPTIVE`: Adaptively switch between 64 and 128 bit floating point numbers.
"""
@enum PrecisionOption begin
    PRECISION_FIXED_64
    PRECISION_FIXED_128
    PRECISION_ADAPTIVE
end

"""
    CoreTrackerOptions

The set of options set for a [`CoreTracker`](@ref). See the description of [`CoreTracker`](@ref)
for all possible options.
"""
mutable struct CoreTrackerOptions
    accuracy::Float64
    initial_step_size::Union{Nothing,Float64}
    max_cond::Float64
    max_corrector_iters::Int
    max_step_size::Float64
    max_steps::Int
    min_step_size::Float64
    precision::PrecisionOption
    simple_step_size_alg::Bool
    steps_jacobian_info_update::Int
    terminate_ill_conditioned::Bool
    track_cond::Bool
    update_patch::Bool
    # Not changeable options
    logarithmic_time_scale::Bool
    from_infinity::Bool
end

function CoreTrackerOptions(;
    # internal
    parameter_homotopy = false,
    update_patch = true,
    logarithmic_time_scale = false,
    from_infinity = false,
    # public
    accuracy = 1e-7,
    initial_step_size = nothing,
    max_cond::Float64 = 1e12,
    max_corrector_iters::Int = 2,
    max_step_size = Inf,
    max_steps = parameter_homotopy ? 10_000 : 1_000,
    min_step_size = 1e-14,
    precision::Union{PrecisionOption,Symbol} = :double,
    simple_step_size_alg = false,
    steps_jacobian_info_update::Int = 5,
    terminate_ill_conditioned::Bool = true,
    track_cond::Bool = true,
     # deprecated
    max_lost_digits = nothing,
    max_refinement_iters = nothing,
    refinement_accuracy = nothing,
)

    max_lost_digits !== nothing && @warn( "Passing `max_lost_digits` to `CoreTracker` is deprecated.",)
    refinement_accuracy !== nothing && @warn(
        "Passing `refinement_accuracy` to `CoreTracker` is deprecated. " *
        "Solutions are now automatically refined to the maximal achievable accuracy.",
    )
    max_refinement_iters !== nothing && @warn(
        "Passing `max_refinement_iters` to `CoreTracker` is deprecated. " *
        "Solutions are now automatically refined to the maximal achievable accuracy.",
    )

    CoreTrackerOptions(
        accuracy,
        initial_step_size,
        max_cond,
        max_corrector_iters,
        max_step_size,
        max_steps,
        min_step_size,
        make_precision(precision),
        simple_step_size_alg,
        steps_jacobian_info_update,
        terminate_ill_conditioned,
        track_cond,
        update_patch,
        logarithmic_time_scale,
        from_infinity,
    )
end

make_precision(p::PrecisionOption) = p
function make_precision(p::Symbol)
    if p == :double
        PRECISION_FIXED_64
    elseif p == :double_double
        PRECISION_FIXED_128
    elseif p == :adaptive
        PRECISION_ADAPTIVE
    else
        throw(ArgumentError(
            "Unsupported argument `precision=$p`. " *
            "Possible values are `:double`, `:double_double` and `:adaptive`.",
        ))
    end
end

Base.show(io::IO, opts::CoreTrackerOptions) = print_fieldnames(io, opts)
Base.show(io::IO, ::MIME"application/prs.juno.inline", opts::CoreTrackerOptions) = opts

Base.copy(opts::CoreTrackerOptions) = copy!(CoreTrackerOptions(), opts)
function Base.copy!(opts::CoreTrackerOptions, opts2::CoreTrackerOptions)
    @unpack accuracy,
    initial_step_size,
    min_step_size,
    max_cond,
    max_corrector_iters,
    max_step_size,
    max_steps,
    precision,
    simple_step_size_alg,
    steps_jacobian_info_update,
    terminate_ill_conditioned,
    track_cond,
    update_patch,
    logarithmic_time_scale,
    from_infinity = opts2
    @pack! opts = accuracy,
    initial_step_size,
    min_step_size,
    max_cond,
    max_corrector_iters,
    max_step_size,
    max_steps,
    precision,
    simple_step_size_alg,
    steps_jacobian_info_update,
    terminate_ill_conditioned,
    track_cond,
    update_patch,
    logarithmic_time_scale,
    from_infinity
    opts
end

####################
# CoreTrackerState #
####################

mutable struct CoreTrackerState{
    T,
    AV<:AbstractVector{Complex{T}},
    AN<:AbstractNorm,
    MaybePatchState<:Union{AbstractAffinePatchState,Nothing},
}
    x::AV # current x
    x̂::AV # last prediction
    x̄::AV # candidate for new x
    ẋ::Vector{Complex{T}} # derivative at current x

    segment::ComplexSegment
    s::Float64 # current step length (0 ≤ s ≤ length(segment))
    Δs::Float64 # current step size
    Δs_prev::Float64 # previous step size

    accuracy::Float64 # norm(x - x(t))
    norm_Δx₀::Float64 # norm of the first newton update
    ω::Float64 # liptschitz constant estimate, see arxiv:1902.02968
    limit_accuracy::Float64
    eval_err::Float64

    norm::AN
    jacobian::JacobianMonitor{T}
    residual::Vector{Float64}

    steps_jacobian_info_update::Int
    status::CoreTrackerStatus.states
    patch::MaybePatchState

    accepted_steps::Int
    rejected_steps::Int
    last_step_failed::Bool
    consecutive_successfull_steps::Int
end

function CoreTrackerState(
    H,
    x₁::AbstractVector,
    t₁,
    t₀,
    options::CoreTrackerOptions,
    patch::Union{Nothing,AbstractAffinePatchState},
    norm::AbstractNorm,
)
    x = isa(x₁, SVector) ? Vector(x₁) : copy(x₁)
    x̂, x̄ = copy(x), copy(x)
    ẋ = Vector(x)
    segment = ComplexSegment(promote(complex(t₁), complex(t₀))...)
    s = 0.0
    Δs = convert(
        Float64,
        min(unpack(options.initial_step_size, Inf), length(segment), options.max_step_size),
    )
    Δs_prev = 0.0
    norm_Δx₀ = accuracy = 0.0
    limit_accuracy = eval_err = eps()
    ω = 1.0
    JM = JacobianMonitor(jacobian(H, x, t₁))
    residual = zeros(first(size(H)))
    steps_jacobian_info_update = 0
    accepted_steps = rejected_steps = 0
    status = CoreTrackerStatus.tracking
    last_step_failed = true
    consecutive_successfull_steps = 0
    CoreTrackerState(
        x,
        x̂,
        x̄,
        ẋ,
        segment,
        s,
        Δs,
        Δs_prev,
        accuracy,
        norm_Δx₀,
        ω,
        limit_accuracy,
        eval_err,
        norm,
        JM,
        residual,
        steps_jacobian_info_update,
        status,
        patch,
        accepted_steps,
        rejected_steps,
        last_step_failed,
        consecutive_successfull_steps,
    )
end

Base.show(io::IO, state::CoreTrackerState) = print_fieldnames(io, state)
Base.show(io::IO, ::MIME"application/prs.juno.inline", state::CoreTrackerState) = state

##################
## CoreTracker ##
##################
"""
    CoreTracker

A `CoreTracker` is the low-level path tracker. Its job is to track an initial solution
`x₁` from `t₁` to `t₀` by following a solution path ``x(t)``. For this a *predictor-corrector* scheme is used.
See our [introduction guide](https://www.juliahomotopycontinuation.org/guides/introduction/#tracking-solution-paths)
for a high-level explanation of the predictor corrector scheme.
The `CoreTracker` accepts for a value `t` a solution ``x^*`` if ``||x(t) - x^*|| < τ`` where ``τ``
is controlled by the `accuracy` option.

To interact with a `CoreTracker` take a look at [`track`](@ref) resp. [`track!`](@ref).
You can also use `CoreTracker` as an iterator. For this you need to call
[`init!(::CoreTracker, x₁, t₁, t₀)`](@ref) first. There is also [`iterator`](@ref) if you
are only interested in pairs `(x,t)`.

The `CoreTracker` **cannot** handle singular solutions or divergent paths. For this you have
to use a [`PathTracker`](@ref) (which uses internally again a `CoreTracker`).

The `CoreTracker` can track solutions in projective space if the underlying homotopy is homogeneous
and the provided start solution is a projective vector (with type [`PVector`](https://www.juliahomotopycontinuation.org/ProjectiveVectors.jl/stable/#ProjectiveVectors.PVector))
from the [`ProjectiveVectors.jl`](https://github.com/JuliaHomotopyContinuation/ProjectiveVectors.jl) package.

The `CoreTracker` can also perform a reparamterization of the parameter ``t`` from ``[0,1]`` towards
``[0,∞]`` by ``t = e^{-s}``. If the homotopy assumes ``t ∈ [0,1]`` this can be done by
passing the `log_transform = true` option. If you already have a homotopy which assumes
``s ∈ [0,∞]`` you can use the `logarithmic_time_scale = true` to tell the core tracker this.
In this case the `min_step_size` criterion is still applied to ``t`` and not ``s``.

## Options

`CoreTracker` accepts many different options. Here are the options in alphabetic order:

* `accuracy` (default `1e-7`): The maximal acceptable distance to the true solution at a
  time step `t`.
* `auto_scaling` (default `true`): Enables an automatic scaling of the variables by
  introducing a *weighted norm*. See also [`WeightedNorm`](@ref).
* `initial_step_size`: The size of the first step. By default an automatic estimate is
  computed but this can be overwritten with this option.
* `max_cond` (default `1e10`): The maximal condition number before a path is terminated due
  to ill-conditioning. This is only applicable if `terminate_ill_conditioned` is `true`.
* `max_corrector_iters (default `2`): The maximal number of Newton iteration steps until
  which the desired accuracy has to be achieved.
* `max_step_size` (default `Inf`): Limit the path tracker to a maximal step size.
* `max_steps` (default `1000`): The maximal number of steps the path tracker makes.
  Note that this changes to `10_000` for parameter homotopies.
* `min_step_size` (default `1e-14`): The path tracking is terminated if the step size is
  smaller than this value.
* `norm` (default [`InfNorm()`](@ref)): The norm used in the computations. Note that if the
  `auto_scaling` is true, this will automatically be converted to a [`WeightedNorm`](@ref).
* `patch`: It is possible to perform tracking in projective space. For the actual
  compuations we still need to choose an affine chart. There are multiple charts possible,
  see [`AbstractAffinePatch`](@ref). The default is [`RandomPatch`](@ref).
* `precision` (default `:double`): This controls the precision used in the computation
  of ``H(x,t)``. If it is set to `:double` only double precision arithmetic is used
  (`Complex{Float64}`).  If it set to `:double_double` double double arithmetic
  as implemented by the [`DoubleFloats`](https://github.com/JuliaMath/DoubleFloats.jl)
  package is **always** used. The option `:adaptive` changes adaptively the precision.
* `predictor` (default [`Pade21`](@ref)): The predictor used.
* `simple_step_size_alg` (default `false`): By default the step size algorithm presented in
  [arXiv:1902.02968](https://arxiv.org/abs/1902.02968) is used. If this option is true
  a more simple step size algorithm is used which doubles the step size after 5 consecutive
  successive steps and halfes the step size after a failed step.
* `terminate_ill_conditioned` (default `true`): Indicates whether the path tracking should be
  terminated for ill-conditioned paths. A path is considered ill-conditioned if the desirable
  accuracy is no more achievable or if the condition number of the Jacobian is larger than
  `1e14`.

"""
struct CoreTracker{
    T,
    AV<:AbstractVector{Complex{T}},
    AN<:AbstractNorm,
    P<:AbstractPredictorCache,
    AP<:Union{Nothing,AbstractAffinePatchState},
    H<:HomotopyWithCache,
    NCC<:NewtonCorrectorCache,
}
    homotopy::H
    predictor::P
    corrector::NCC
    # these are mutable
    state::CoreTrackerState{T,AV,AN,AP}
    options::CoreTrackerOptions
    # buffer
    dt::Vector{Complex{T}}
end

function CoreTracker(
    prob::AbstractProblem,
    x₁;
    log_transform::Bool = false,
    log_homotopy::Bool = log_transform,
    logarithmic_time_scale = log_transform || log_homotopy,
    kwargs...,
)
    if log_homotopy
        H = LogHomotopy(prob.homotopy)
        CoreTracker(
            H,
            embed(prob, x₁),
            complex(0.0),
            complex(36.0),
            prob;
            logarithmic_time_scale = logarithmic_time_scale,
            kwargs...,
        )
    else
        CoreTracker(
            prob.homotopy,
            embed(prob, x₁),
            complex(1.0),
            complex(0.0),
            prob;
            logarithmic_time_scale = logarithmic_time_scale,
            kwargs...,
        )
    end
end

# Tracking in affine space
function CoreTracker(
    homotopy::AbstractHomotopy,
    x₁::AbstractVector,
    t₁::Number,
    t₀::Number,
    prob::AbstractProblem;
    patch = nothing,
    corrector::NewtonCorrector = NewtonCorrector(),
    predictor::AbstractPredictor = Pade21(),
    auto_scaling::Bool = true,
    norm::AbstractNorm = InfNorm(),
    kwargs...,
)

    patch === nothing || throw(ArgumentError("You can only pass `patch=$(patch)` if `affine_tracking=false`."))

    options = CoreTrackerOptions(;
        parameter_homotopy = isa(homotopy, ParameterHomotopy), kwargs...,
    )

    # We close over the homotopy and its cache to be able to pass things around more easily
    H = HomotopyWithCache(homotopy, Vector(x₁), t₁)

    used_norm = auto_scaling ? WeightedNorm(norm, x₁) : norm
    # We have to make sure that the element type of x is invariant under evaluation
    state =
        CoreTrackerState(H, indempotent_x(H, x₁, t₁), t₁, t₀, options, nothing, used_norm)

    pred_cache = cache(predictor, H, state.x, state.ẋ, t₁)
    corr_cache = cache(corrector, H, state.x, t₁)
    dt = Vector{eltype(state.ẋ)}(undef, first(size(H)))

    CoreTracker(H, pred_cache, corr_cache, state, options, dt)
end

# Tracking in Projective Space
function CoreTracker(
    homotopy::AbstractHomotopy,
    x₁::ProjectiveVectors.PVector,
    t₁::Number,
    t₀::Number,
    prob::AbstractProblem;
    patch = RandomPatch(),
    corrector::NewtonCorrector = NewtonCorrector(),
    predictor::AbstractPredictor = Pade21(),
    auto_scaling::Bool = true,
    norm::AbstractNorm = InfNorm(),
    simple_step_size_alg = isa(patch, OrthogonalPatch),
    kwargs...,
)

    options = CoreTrackerOptions(;
        parameter_homotopy = isa(homotopy, ParameterHomotopy), kwargs...,
    )

    if homotopy isa PatchedHomotopy
        error("You cannot pass a `PatchedHomotopy` to CoreTracker. Instead pass the homotopy and patch separate.")
    end

    patch_state = state(patch, x₁)

    # We close over the patch state, the homotopy and its cache
    # to be able to pass things around more easily
    H = HomotopyWithCache(PatchedHomotopy(homotopy, patch_state), x₁, t₁)

    if is_global_patch(patch) && auto_scaling
        used_norm = WeightedNorm(norm, x₁)
    else
        used_norm = norm
    end
    # We have to make sure that the element type of x is invariant under evaluation
    ct_state = CoreTrackerState(
        H,
        indempotent_x(H, x₁, t₁),
        t₁,
        t₀,
        options,
        patch_state,
        used_norm,
    )

    pred_cache = cache(predictor, H, ct_state.x, ct_state.ẋ, t₁)
    corr_cache = cache(corrector, H, ct_state.x, t₁)
    dt = Vector{eltype(ct_state.ẋ)}(undef, first(size(H)))

    CoreTracker(H, pred_cache, corr_cache, ct_state, options, dt)
end

"""
    indempotent_x(H, x₁, t₁)

This returns a vector similar to `x₁` but with an element type which is invariant under evaluation.
"""
function indempotent_x(H::AbstractHomotopy, x₁, t₁)
    u = Vector{Any}(undef, size(H)[1])
    evaluate!(u, H, x₁, t₁)

    if isa(x₁, SVector)
        indem_x = Vector{promote_type(typeof(u[1]), ComplexF64)}(undef, length(x₁))
    else
        indem_x = similar(x₁, promote_type(typeof(u[1]), ComplexF64))
    end
    indem_x .= x₁
    indem_x
end

Base.show(io::IO, C::CoreTracker) =
    print(io, "CoreTracker tracking a path of type $(typeof(C.state.x))")
Base.show(io::IO, ::MIME"application/prs.juno.inline", x::CoreTracker) = x

Base.broadcastable(C::CoreTracker) = Ref(C)

###########
## Alias ##
###########
const CT = CoreTracker
const CTS = CoreTrackerState
const CTO = CoreTrackerOptions
const CTR = CoreTrackerResult

#######################
## Correct & Tangent ##
#######################

function compute_ẋ!(tracker::CT; update_jacobian::Bool = true)
    @unpack homotopy, state, options, dt = tracker

    if update_jacobian
        jacobian_and_dt!(jacobian(state.jacobian), dt, homotopy, state.x, current_t(state))
        updated!(state.jacobian)
    else
        dt!(dt, homotopy, state.x, current_t(state))
    end
    # flip sign of entries in dt
    @inbounds for i in eachindex(dt)
        dt[i] = -dt[i]
    end
    # solve the linear system to compute ẋ
    LA.ldiv!(state.ẋ, state.jacobian, dt)
    state.ẋ
end

@inline function correct!(
    x̄::AbstractVector,
    tracker::CT,
    x::AbstractVector = tracker.state.x,
    t::Number = current_t(tracker.state);
    simplified_last_step::Bool = true,
    min_iters::Int = 1,
    max_iters::Int = tracker.options.max_corrector_iters + 1,
    full_steps::Int = max_iters - simplified_last_step,
    tol::Float64 = tracker.options.accuracy,

)
    if tracker.options.precision == PRECISION_ADAPTIVE
        double_64_evaluation =
            at_limit_accuracy(tracker.state, tracker.options; safety_factor = 100.0)
    else
        double_64_evaluation = tracker.options.precision == PRECISION_FIXED_128
    end

    newton!(
        x̄,
        tracker.homotopy,
        x,
        t,
        tracker.state.jacobian,
        tracker.state.norm,
        tracker.corrector;
        tol = tol,
        max_iters = max_iters,
        min_iters = min_iters,
        double_64_evaluation = double_64_evaluation,
        full_steps = full_steps,
    )
end

function limit_accuracy!(
    tracker::CT,
    x::AbstractVector = tracker.state.x,
    t::Number = current_t(tracker.state),
    norm::AbstractNorm = tracker.state.norm;
    update_cond::Bool = tracker.options.track_cond,
    update_eval_err::Bool = true,
)
    @unpack state, options = tracker

    for i = 1:2
        # We do a second iteration if we would terminate the path tracking to make sure
        # that this is not just an artifact
        i == 2 && !terminate_limit_accuracy(state, options) && continue
        limit_acc, eval_err = limit_accuracy!(
            tracker.homotopy,
            x,
            t,
            state.jacobian,
            norm,
            tracker.corrector;
            x_accuracy = state.accuracy,
            update_jacobian = false,
            compute_cond = update_cond,
            compute_eval_err = update_eval_err,
        )
        state.accuracy = min(state.accuracy, limit_acc)
        state.limit_accuracy = max(state.accuracy, limit_acc)
        if update_eval_err
            state.eval_err = exp((log(state.eval_err) + log(eval_err)) / 2)
        end
    end

    nothing
end

function terminate_limit_accuracy(state, options)
    (
        options.precision == PRECISION_FIXED_64
    ) && at_limit_accuracy(state, options; safety_factor = 10.0)
end
function at_limit_accuracy(state::CTS, opts::CTO; safety_factor::Float64 = 10.0)
    opts.accuracy < safety_factor * state.limit_accuracy
end
####################
## INITIALIZATION ##
####################

embed!(x::ProjectiveVectors.PVector, y) = ProjectiveVectors.embed!(x, y)
embed!(x::AbstractVector, y) = x .= y

function check_start_value!(tracker::CT, x = tracker.state.x, t = current_t(tracker.state))
    # We perturb the provided soution by sligthly more than the desired accuracy
    # to check that the Newton iteration converges properly and to obtain an estimate
    # of ω
    perturb!(tracker.state.x̄, x, 10 * tracker.options.accuracy, tracker.state.norm)

    # Solutions could need higher order precision already at the beginning
    # Therefore, if the standard check fails, try again with higher precision (if applicable)
    double_64_evaluation = false
    for i = 1:2
        result = newton!(
            tracker.state.x̄,
            tracker.homotopy,
            tracker.state.x̄,
            t,
            tracker.state.jacobian,
            tracker.state.norm,
            tracker.corrector;
            tol = tracker.options.accuracy,
            max_iters = 2,
            simplified_last_step = false,
            double_64_evaluation = double_64_evaluation,
        )
        if is_converged(result)
            tracker.state.ω = result.ω₀
            x .= tracker.state.x̄
            tracker.state.accuracy = result.accuracy
            limit_accuracy!(tracker)
            break
        elseif i == 1 && tracker.options.precision != PRECISION_FIXED_64
            double_64_evaluation = true
        else
            tracker.state.status = CoreTrackerStatus.terminated_invalid_startvalue
            break
        end
    end
    nothing
end

perturb!(y, x, Δ::Real, w::WeightedNorm) = y .= x .+ Δ .* weights(w)
perturb!(y, x, Δ::Real, norm::AbstractNorm) = y .= x .+ Δ

"""
    init!(tracker::CT, x₁, t₁, t₀;
      setup_patch::Bool = tracker.options.update_patch,
      loop::Bool = false,
      check_start_value::Bool = !loop)

Setup `tracker` to track `x₁` from `t₁` to `t₀`. Use this if you want to use the
`tracker` as an iterator.
"""
function init!(
    tracker::CoreTracker,
    x₁::AbstractVector,
    t₁,
    t₀;
    setup_patch::Bool = tracker.options.update_patch,
    loop::Bool = false,
    keep_steps::Bool = false,
    check_start_value::Bool = !loop,
)
    @unpack state, predictor, homotopy = tracker

    init!(state, x₁, t₁, t₀, tracker.options, setup_patch, loop, keep_steps)
    if !loop
        check_start_value!(tracker)
        compute_ẋ!(tracker; update_jacobian = false)
        init!(predictor, homotopy, state.x, state.ẋ, current_t(state), state.jacobian)
    end
    initial_step_size!(state, predictor, tracker.options)
    tracker
end

function init!(
    state::CTS,
    x₁::AbstractVector,
    t₁,
    t₀,
    options::CTO,
    setup_patch::Bool,
    loop::Bool,
    keep_steps::Bool,
)
    state.segment = ComplexSegment(t₁, t₀)
    state.s = 0.0
    state.Δs_prev = 0.0
    state.status = CoreTrackerStatus.tracking
    embed!(state.x, x₁)
    setup_patch && state.patch !== nothing && init!(state.patch, state.x)
    if !loop
        state.consecutive_successfull_steps = 0
        state.last_step_failed = true
        state.accuracy = 0.0
        state.ω = 1.0
        state.eval_err = state.limit_accuracy = eps()
        if !keep_steps
            state.accepted_steps = state.rejected_steps = 0
        end
        init!(state.norm, state.x)
        init!(state.jacobian)
        state.residual .= 0.0
        state.steps_jacobian_info_update = 0
    end
    state
end

@deprecate setup!(tracker::CoreTracker, x₁, t₁, t₀; kwargs...) init!(
    tracker,
    x₁,
    t₁,
    t₀;
    kwargs...,
)

##############
## TRACKING ##
##############

"""
    track(tracker::CoreTracker, x₁, t₁=1.0, t₀=0.0)::CoreTrackerResult

Track a value `x₁` from `t₁` to `t₀` using the given `CoreTracker` `tracker`.
This returns a `CoreTrackerResult`. This modifies `tracker`.
See [`track!`](@ref) for the possible options.
To investigate the behaviour of a particular take a look at [`path_info`](@ref).
"""
function track(
    tracker::CT,
    x₁::AbstractVector,
    t₁::Number = 1.0,
    t₀::Number = 0.0;
    kwargs...,
)
    track!(tracker, x₁, t₁, t₀; kwargs...)
    CoreTrackerResult(tracker)
end

"""
    track!(tracker::CoreTracker, x₁, t₁=1.0, t₀=0.0)::CoreTrackerStatus.states

Track a value `x₁` from `t₁` to `t₀` using the given `CoreTracker` `tracker`.
Returns one of the enum values of `CoreTrackerStatus` indicating the status.
Check [`is_success`](@ref) and [`is_terminated`](@ref) to test for the status.

If `setup_patch` is `true` then [`init!`](@ref) is called at the beginning
of the tracking.

    track!(x₀, tracker::CoreTracker, x₁, t₁=1.0, t₀=0.0)::CoreTrackerStatus.states

Additionally also stores the result in `x₀` if the tracking was successfull.
"""
function track!(
    tracker::CT,
    x₁::AbstractVector,
    t₁::Number = 1.0,
    t₀::Number = 0.0;
    setup_patch::Bool = tracker.options.update_patch,
    loop::Bool = false,
    check_start_value::Bool = !loop,
    debug::Bool = false,
)
    _track!(tracker, x₁, t₁, t₀, setup_patch, loop, check_start_value, debug)
end

function track!(
    x₀::AbstractVector,
    tracker::CT,
    x₁::AbstractVector,
    t₁::Number = 1.0,
    t₀::Number = 0.0;
    setup_patch::Bool = tracker.options.update_patch,
    loop::Bool = false,
    check_start_value::Bool = !loop,
    debug::Bool = false,
)

    _track!(tracker, x₁, t₁, t₀, setup_patch, loop, check_start_value, debug)
    retcode = status(tracker)
    if is_success(status(tracker))
        x₀ .= current_x(tracker)
    end
    retcode
end
@inline function _track!(
    tracker::CT,
    x₁::AbstractVector,
    t₁::Number,
    t₀::Number,
    setup_patch::Bool,
    loop::Bool,
    check_start_value::Bool,
    debug::Bool,
)

    # no need to track, so just assign solution and we are done
    if t₁ == t₀
        embed!(tracker.state.x, x₁)
        return (tracker.state.status = CoreTrackerStatus.success)
    end

    init!(
        tracker,
        x₁,
        t₁,
        t₀;
        setup_patch = setup_patch,
        check_start_value = check_start_value,
        loop = loop,
    )

    while is_tracking(tracker.state.status)
        step!(tracker, debug)
    end

    tracker.state.status
end

"""
    step!(tracker::CT)

Perform one step of the path tracker. This tries to move forward by `Δt`, i.e.,
obtaining an approximation of `x(t+Δt)`.
Returns `true` if the step was successfull.
"""
function step!(tracker::CT, debug::Bool = false)
    @unpack homotopy, predictor, state, options = tracker
    @unpack x, x̂, x̄, ẋ = state

    # @unpack currently doesn't work with getproperty
    t, Δt = current_t(state), current_Δt(state)

    debug && printstyled(
        "\nt: ",
        round(t, sigdigits = 5),
        " Δt: ",
        round(Δt; sigdigits = 5),
        "\n";
        color = :yellow,
    )

    # Use the current approximation of x(t) to obtain estimate
    # x̂ ≈ x(t + Δt) using the provided predictor
    predict!(x̂, predictor, homotopy, x, t, Δt, ẋ, state.jacobian)
    # Correct the predicted value x̂ to obtain x̄.
    # If the correction is successfull we have x̄ ≈ x(t+Δt).
    result = correct!(x̄, tracker, x̂, t + Δt; simplified_last_step = true)

    debug && show(result)
    if is_converged(result)
        # move forward
        x .= x̄
        state.s += state.Δs
        state.accuracy = result.accuracy

        # update step size
        state.Δs_prev = state.Δs
        update_stepsize!(tracker, result)

        # If we use dynamic affine charts, now we update to a new one
        if state.patch !== nothing && options.update_patch
            changepatch!(state.patch, x)
        end

        # If we use a weighted norm, now we update this
        update!(state.norm, x)

        # Compute the new tangent ẋ(t).
        compute_ẋ!(tracker; update_jacobian = true)

        # Make an additional newton step to check the limiting accuracy
        if (state.steps_jacobian_info_update += 1) ≥ options.steps_jacobian_info_update ||
           state.last_step_failed
            limit_accuracy!(tracker)
            state.steps_jacobian_info_update = 0
        end
        # tell the predictors about the new derivative if they need to update something
        update!(predictor, homotopy, x, ẋ, t + Δt, state.jacobian, state.eval_err)

        # Update other state
        state.accepted_steps += 1
        state.last_step_failed = false
    else
        # Step failed, so we have to try with a new (smaller) step size
        update_stepsize!(tracker, result)

        # Update other state
        state.rejected_steps += 1
        state.last_step_failed = true
    end
    state.norm_Δx₀ = result.norm_Δx₀

    check_terminated!(state, options)

    # If we terminate, always update the limit accuracy (and by this refine the solution
    # to maximal accuracy).
    if is_success(state.status) && state.steps_jacobian_info_update ≠ 0
        limit_accuracy!(tracker)
    end

    state.last_step_failed
end

###############
## STEP SIZE ##
###############

## intial step size
function initial_step_size!(state::CTS, predictor::AbstractPredictorCache, options::CTO)
    if options.initial_step_size !== nothing
        return state.Δs = max(
            min(options.initial_step_size, length(state.segment), options.max_step_size),
            sqrt(options.min_step_size),
        )
    end
    ω = max(state.ω, 1e-8)
    p = order(predictor)
    deriv, k = unpack(highest_derivative(predictor), (state.ẋ, 1))
    η_p = nthroot(state.norm(deriv) / factorial(k), k)
    δ_N_ω = δ(options, ω, options.max_corrector_iters / 2)
    Δs = nthroot(g(δ_N_ω) / ω, p) / η_p
    if isnan(Δs)
        Δs = 0.05 * length(state.segment)
    end

    state.Δs =
        max(min(Δs, length(state.segment), options.max_step_size), options.min_step_size)
end

## Step size update

function update_stepsize!(tracker::CT, result::NewtonCorrectorResult)
    @unpack state, options = tracker

    if !isnan(result.ω₀)
        # Use result.ω₀ in the converged case to avoid wrong values of ω due to
        # hitting the accuracy limit.
        if is_converged(result) || isnan(state.ω)
            state.ω = result.ω₀
        else
            state.ω = max(result.ω₀, state.ω)
        end
    end

    # Use the simple step size algorithm in the case that we get very close to
    # the accuracy limit. The motivation is that we can run into danger that the computed
    # estimate of ω is too high due to hitting limit accuracy.
    near_accuracy_limit = at_limit_accuracy(state, options; safety_factor = 100.0)
    if options.simple_step_size_alg || (
        near_accuracy_limit && options.precision == PRECISION_FIXED_64
    )
        Δs = simple_step_size_alg!(state, options, result)
    else
        Δs = adaptive_step_size_alg!(state, options, result, order(tracker.predictor))
    end

    # Make sure to not overshoot.
    if !is_min_step_size(Δs, state.s, length(state.segment), options) && steps(state) > 5
        state.status = CoreTrackerStatus.terminated_step_size_too_small
    else
        state.Δs = min(length(state.segment) - state.s, Δs, options.max_step_size)
    end

    nothing
end

function is_min_step_size(Δs, s, length_segment, options::CTO)
    if options.logarithmic_time_scale && options.from_infinity
        ŝ = length_segment - s
        t = exp(-ŝ)
        Δt = exp(-ŝ + Δs) - t
        Δt ≥ min(options.min_step_size, t / 100)
    elseif options.logarithmic_time_scale
        exp(-s) - exp(-(s + Δs)) ≥ options.min_step_size
    else
        Δs ≥ options.min_step_size
    end
end

## Simple step size algorithm

function simple_step_size_alg!(state::CTS, options::CTO, result::NewtonCorrectorResult)

    if !is_converged(result)
        state.consecutive_successfull_steps = 0
        return state.Δs / 2
    end

    state.consecutive_successfull_steps = (state.consecutive_successfull_steps + 1) % 5
    state.consecutive_successfull_steps == 0 ? 2 * state.Δs : state.Δs
end

## Adaptive step size algorithm

"Use the adaptive step size control described in https://arxiv.org/abs/1902.02968"
function adaptive_step_size_alg!(
    state::CTS,
    options::CTO,
    result::NewtonCorrectorResult,
    ord::Int,
)
    @unpack s, ω = state

    d_x̂_x = result.iters > 1 ? state.norm(state.x̂, state.x) : result.norm_Δx₀
    Δx₀ = result.norm_Δx₀

    if is_converged(result)
        # This is to handle the edge case that g(δ_N_ω) >> (ω * d_x̂_x̄) > 0 but
        # at the same time δ_N_ω < eps(). Since then g(δ_N_ω) = 0
        δ_N_ω = max(δ(options, ω, options.max_corrector_iters / 4), 1e-15)
        Δs = nthroot(g(δ_N_ω) / (ω * d_x̂_x), ord) * state.Δs
        if state.last_step_failed
            Δs = min(Δs, state.Δs)
        end
    else
        δ_N_ω = δ(options, ω, options.max_corrector_iters / 4)
        ω_η = 0.5ω * Δx₀
        if δ_N_ω < ω_η
            Δs = min(nthroot(g(δ_N_ω) / g(ω_η), ord), 0.9) * state.Δs
        else
            Δs = 0.5 * state.Δs
        end
    end

    Δs
end

g(Θ::Real) = sqrt(1.0 + 4Θ) - 1.0
δ(opts::CTO, ω::Real, μ::Real) = min(√(0.5ω) * τ(opts, μ), 0.25)
function τ(opts::CTO, μ::Real)
    # most common case: 2(2 - 0.5) == 3
    if opts.max_corrector_iters == 2 && μ == 0.5
        cbrt(opts.accuracy)
    else
        opts.accuracy^(1 / (2 * (opts.max_corrector_iters - μ)))
    end
end

function check_terminated!(state::CTS, options::CTO)
    Δs = length(state.segment) - state.s
    if Δs < options.min_step_size
        state.status = CoreTrackerStatus.success
        state.s = length(state.segment)
    elseif steps(state) ≥ options.max_steps
        state.status = CoreTrackerStatus.terminated_maximal_iterations
    elseif options.terminate_ill_conditioned && terminate_limit_accuracy(state, options)
        state.status = CoreTrackerStatus.terminated_accuracy_limit
    elseif options.terminate_ill_conditioned && cond(state) > options.max_cond
        state.status = CoreTrackerStatus.terminated_ill_conditioned
    end
    nothing
end

function Base.iterate(tracker::CoreTracker, state::Int = 0)
    if is_tracking(tracker.state.status)
        step!(tracker)
        tracker, state + 1
    else
        nothing
    end
end

function is_valid_start_value(tracker::CoreTracker, x::AbstractVector, t::Number)
    @unpack state = tracker
    embed!(state.x̄, x)
    state.patch !== nothing && init!(state.patch, state.x̄)
    check_start_value!(tracker, state.x̄, t)
    !is_invalid_startvalue(status(tracker))
end

"""
    affine_tracking(tracker)

Returns `true` if the path tracker tracks in affine space. If `false` then then the path tracker
tracks in some affine chart of the projective space.
"""
affine_tracking(tracker::CT) = !isa(tracker.state.x, ProjectiveVectors.PVector)

#################
## Query State ##
#################
"""
    current_t(tracker::CoreTracker)

Current `t`.
"""
current_t(tracker::CoreTracker) = current_t(tracker.state)
current_t(state::CTS) = state.segment[state.s]

"""
    current_Δt(tracker::CoreTracker)

Current step_size `Δt`.
"""
current_Δt(tracker::CoreTracker) = current_Δt(tracker.state)
current_Δt(state::CTS) = state.segment[state.Δs] - state.segment.start

"""
    steps(tracker::CoreTracker)

Current number of steps.
"""
steps(tracker::CoreTracker) = steps(tracker.state)
steps(state::CTS) = state.accepted_steps + state.rejected_steps

@deprecate iters(tracker::CoreTracker) steps(tracker)

"""
    status(tracker::CoreTracker)

Current status.
"""
status(tracker::CoreTracker) = status(tracker.state)
status(state::CTS) = state.status

"""
    current_x(tracker::CoreTracker)

Return the current value of `x`.
"""
current_x(tracker::CoreTracker) = current_x(tracker.state)
current_x(state::CTS) = state.x

"""
    LinearAlgebra.cond(tracker::CoreTracker)

Returns the currently computed approximation of the condition number of the
Jacobian.
"""
LA.cond(tracker::CoreTracker) = cond(tracker.state)
LA.cond(state::CTS) = cond(state.jacobian)

"""
    norm(tracker::CoreTracker)

Returns the norm used to compute distances during the path tracking.
"""
LA.norm(tracker::CoreTracker) = LA.norm(tracker.state)
LA.norm(state::CTS) = state.norm

"""
    options(tracker::CoreTracker)

Returns the options used in the tracker.
"""
options(tracker::CoreTracker) = tracker.options

##################
# Modify options #
##################
"""
    accuracy(tracker::CoreTracker)

Current accuracy.
"""
accuracy(tracker::CoreTracker) = tracker.options.accuracy

"""
    set_accuracy!(tracker::CoreTracker, accuracy)

Set the current accuracy to `accuracy`.
"""
function set_accuracy!(tracker::CoreTracker, accuracy)
    @unpack options = tracker
    options.accuracy = accuracy
    tracker
end

"""
    max_corrector_iters(tracker::CoreTracker)

Current correction max_steps.
"""
max_corrector_iters(T::CT) = T.options.max_corrector_iters

"""
 set_max_corrector_iters!(tracker::CoreTracker, n)

Set the correction max_steps to `n`.
"""
set_max_corrector_iters!(T::CT, n) = T.options.max_corrector_iters = n

"""
    max_step_size (tracker::CoreTracker)

Current maximal step size.
"""
max_step_size(T::CT) = T.options.max_step_size

"""
    set_max_corrector_iters!(tracker::CoreTracker, Δs)

Set the maximal step size to `Δs`.
"""
set_max_step_size!(T::CT, Δs) = T.options.max_step_size = Δs

## Deprecated ##
"""
    refinement_accuracy(tracker::CoreTracker)

Current refinement accuracy.
"""
function refinement_accuracy(tracker::CoreTracker)
    @warn("Solutions are now automatically refined to maximal achievable accuracy.")
    return tracker.options.accuracy
end

"""
    set_max_refinement_iters!(tracker::CoreTracker, accuracy)

Set the current refinement accuracy to `accuracy`.
"""
function set_refinement_accuracy!(T::CT, accuracy)
    @warn( "`set_refinement_accuracy!` is deprecated. Solutions are now automatically refined to maximal achievable accuracy.",)
end

"""
    max_refinement_iters(tracker::CoreTracker)

Current refinement max_steps.
"""
function max_refinement_iters(T::CT)
    @warn( "`max_refinement_iters` is deprecated. Solutions are now automatically refined to the maximal achievable accuracy.",)
    T.options.max_corrector_iters
end

"""
    set_max_refinement_iters!(tracker::CoreTracker, n)

Set the current refinement max_steps to `n`.
"""
function set_max_refinement_iters!(T::CT, n)
    @warn( "`set_max_refinement_iters` is deprecated. Solutions are now automatically refined to the maximal achievable accuracy.",)
end

"""
    start_parameters!(tracker::CoreTracker, p)

Set the start parameters of the homotopy in in `tracker` to `p`.
"""
start_parameters!(tracker::CoreTracker, p::AbstractVector) =
    (set_start_parameters!(basehomotopy(tracker.homotopy), p); tracker)

"""
    target_parameters!(tracker::CoreTracker, p)

Set the target parameters of the homotopy in in `tracker` to `p`.
"""
target_parameters!(tracker::CoreTracker, p::AbstractVector) =
    (set_target_parameters!(basehomotopy(tracker.homotopy), p); tracker)

################
# PathIterator #
################
struct PathIterator{Tracker<:CoreTracker}
    tracker::Tracker
    t_real::Bool
end
Base.IteratorSize(::Type{<:PathIterator}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{<:PathIterator}) = Base.HasEltype()

"""
    iterator(tracker::CoreTracker, x₁, t₁=1.0, t₀=0.0; affine=true)

Prepare a tracker to make it usable as a (stateful) iterator. Use this if you want to inspect a specific
path. In each iteration the tuple `(x,t)` is returned.
If `affine == true` then `x` is the affine solution (internally we compute in projective space).

## Example

Assume you have `CoreTracker` `tracker` and you wan to track `x₁` from 1.0 to 0.25:
```julia
for (x,t) in iterator(tracker, x₁, 1.0, 0.25)
println("x at t=\$t:")
println(x)
end
```

Note that this is a stateful iterator. You can still introspect the state of the tracker.
For example to check whether the tracker was successfull
(and did not terminate early due to some problem) you can do
```julia
println("Success: ", status(tracker) == CoreTrackerStatus.success)
```
"""
function iterator(tracker::CT, x₁, t₁ = 1.0, t₀ = 0.0; kwargs...)
    init!(tracker, x₁, t₁, t₀; kwargs...)
    PathIterator(tracker, typeof(t₁ - t₀) <: Real)
end

function current_x_t(iter::PathIterator)
    x = current_x(iter.tracker)
    t = current_t(iter.tracker)
    (copy(x), iter.t_real ? real(t) : t)
end

function Base.iterate(iter::PathIterator, state = nothing)
    state === nothing && return current_x_t(iter), 1
    iter.tracker.state.status != CoreTrackerStatus.tracking && return nothing

    t = current_t(iter.tracker)
    step_done = false
    while !step_done && is_tracking(iter.tracker.state.status)
        step!(iter.tracker)
        step_done = current_t(iter.tracker) != t
    end
    current_x_t(iter), state + 1
end

#############################
# Convencience constructors #
#############################
"""
    coretracker_startsolutions(args...; kwargs...)

Construct a [`CoreTracker`](@ref) and `startsolutions` in the same way `solve`
does it. This also takes the same input arguments as `solve`. This is convenient if you want
to investigate single paths.
"""
function coretracker_startsolutions(args...; kwargs...)
    supported, rest = splitkwargs(kwargs, problem_startsolutions_supported_keywords)
    prob, startsolutions = problem_startsolutions(args...; supported...)
    tracker = CoreTracker(prob, start_solution_sample(startsolutions); rest...)
    (tracker = tracker, startsolutions = startsolutions)
end

"""
    coretracker(args...; kwargs...)

Construct a [`CoreTracker`](@ref) in the same way `solve`
does it. This also takes the same input arguments as `solve` with the exception that you do not need to specify startsolutions.
This is convenient if you want to investigate single paths.

## Examples

### Obtain single solution
We want to construct a path tracker to track a parametrized system `f` with parameters `p`
from the parameters `a` to `b`.
```julia
tracker = coretracker(f; parameters = p, start_parameters = a, target_parameters = b)
```
You then can obtain a single solution at `b` by using
```julia
x_b = track(tracker, x_a).x
```

### Trace a path
To trace a path you can use the [`iterator`](@ref) method.

```julia
tracker = coretracker(f; parameters = p, start_parameters = a, target_parameters = b)
for (x, t) in iterator(tracker, x₁)
    @show (x,t)
end
```

If we want to guarantee smooth traces we can limit the maximal step size.
```julia
tracker = coretracker(f; parameters = p, max_step_size = 0.01,
                         start_parameters = a, target_parameters = b)
for (x, t) in iterator(tracker, x₁)
    @show (x,t)
end
```
"""
coretracker(args...; kwargs...) = first(coretracker_startsolutions(args...; kwargs...))


text:
struct VecArray{C,A,N} <: AbstractVec{C,N}
  batch::A
  sizes
end

Base.show(io::IO, x::VecArray) = println(io, "VecArray")

data(xs::VecArray) = xs.batch

VecArray(batch::C...) where {C <: AbstractArray} = let
  batched, sizes = tobatch(batch...)
  VecArray{C, typeof(batched), length(batch)}(batched, sizes)
end

function tobatch(xs::A...) where {T <: ScalarTypes,N,A <: AbstractArray{T,N}}
  if all(y->y==size(xs[1]), size.(xs))
    res = cat(xs..., dims=N+1)
    return (res, size.(xs))
  end
  dims = max.(size.(xs)...)
  batch = zeros(T, dims..., length(xs)) #TODO: this feels wrooooong
  for (i,x) in enumerate(xs)
    setindex!(batch, x, axes(x)..., i)
  end
  (batch, size.(xs))
end

function select(conds::Mask{N}, first_vals::VecArray{C,T,N}, second_vals::VecArray{C,T,N}) where {C,T,N}
  result = zero(first_vals.batch)
  for (i,cond) in enumerate(conds)
    result[axes(result)[1:end-1]..., i] = cond ?
      first_vals.batch[axes(result)[1:end-1]..., i] :
      second_vals.batch[axes(result)[1:end-1]..., i]
  end
  return VecArray{C,T,N}(result, first_vals.sizes)
end

function select(conds::Mask{N}, first_vals::VecArray{C,T,N}, second_vals::C) where {C <: AbstractArray,T,N}
  result = zero(first_vals.batch)
  for (i,cond) in enumerate(conds)
    result[axes(result)[1:end-1]..., i] = cond ?
      first_vals.batch[axes(result)[1:end-1]..., i] :
      second_vals
  end
  return VecArray{C,T,N}(result, first_vals.sizes)
end

function select(conds::Mask{N}, first_vals::C, second_vals::VecArray{C,T,N}) where {C <: AbstractArray,T,N}
  result = zero(second_vals.batch)
  for (i,cond) in enumerate(conds)
    result[axes(result)[1:end-1]..., i] = cond ?
      first_vals :
      second_vals.batch[axes(result)[1:end-1]..., i]
  end
  return VecArray{C,T,N}(result, first_vals.sizes)
end

vect(xs::AbstractArray{T}...) where {T <: ScalarTypes} = VecArray(xs...)
VecArrayOrArrayOrVal{C,A,N} = Union{VecArray{C,A,N}, C, ScalarTypes}
sizes(xs) = xs[findfirst(x->x isa VecArray, xs)].sizes
# element wise ops
for op in :[+, -, &, |, !,
            ==, !=, >, >=, <, <=].args
    @eval begin
        @spmd $op(xs::VecArrayOrArrayOrVal{C,A,N}...) where {C,A,N} =
          VecArray{C,A,N}($op.(data.(xs)...), sizes(xs))
    end
end

@spmd length(x::VecArray{C,A,N}) where {C,A,N} = vect(map(sum, S)...)
@spmd size(x::VecArray{C,A,N}) where {C,A,N} = vect(S...)
#
# VecVectorOrVector{C <: AbstractArray{T,1}, A <: AbstractArray{T,2}, N, S} = Union{VecArray{C,A,N,S}, C}
# @spmd (*)(xs::VecVectorOrVector{C,A,N,S}...) where {C,A,N,S} =
#   VecArray{C,A,N,map(x->,S)}((*)(data.(xs)...))

@spmd (*)(xs::VecArray{C1,A,N}, ys::C2) where {T,C1<:AbstractVector,A,N,C2<:AbstractMatrix{T}} = let
  create_new_sizes(s1,s2) = (s1[1], length(s2) == 1 ? () : s2[2])
  result = (*)(data.((xs,ys))...)
  sizes = map(s->create_new_sizes(s,size(ys)),xs.sizes)
  VecArray{C2,typeof(result),N}(result, sizes)
end
@spmd (*)(xs::C1, ys::VecArray{C2,A,N}) where {T,C1<:AbstractMatrix{T},A,N,C2<:AbstractVector} = let
  create_new_sizes(s1,s2) = (s1[1], length(s2) == 1 ? () : s2[2])
  result = (*)(data.((xs,ys))...)
  sizes = map(s->create_new_sizes(size(xs),s),ys.sizes)
  VecArray{C2,typeof(result),N}(result, sizes)
end

@spmd Base.Broadcast.broadcasted(f,xs::VecArrayOrArrayOrVal...) = Base.Broadcast.broadcasted(f,data.(xs)...)
@spmd Base.Broadcast.materialize(x) = Base.Broadcast.materialize(x)

type_minus_one(xs::C) where {T,N, C<:AbstractArray{T,N}} = AbstractArray{T,N-1}

@spmd reshape(xs::VecArray{C,A,N}, dims) where {C,A,N} = begin
  res = reshape(data(xs), dims..., N)
  VecArray{type_minus_one(res),typeof(res),N}(res, repeat([dims...], N))
end

@spmd reshape(xs::VecArray{C,A,N}, dims...) where {C,A,N} = begin
  res = reshape(data(xs), dims..., N)
  VecArray{type_minus_one(res),typeof(res),N}(res, repeat([dims...], N))
end


text:
using Iris
using Test

@testset "Helmholtz Operators" begin
    @testset "1D" begin
        cavity = Interval(-.5,.5)
        pump = Interval(-.5,0)
        bnd = Boundary(Interval(-1,1),noBL,DirichletBC)
        dom = NondispersiveDomain(cavity,2)
        fdom = DispersiveDomain(pump,TwoLevelSystem(10,.1,2))
        lat = Cartesian(.05)
        latdom = LatticeDomain(bnd,lat)
        sim = Simulation(10,dom,latdom,fdom)
        h = @test_nowarn Helmholtz(sim; m=2)
        d = @test_nowarn h(0)
        @test d[1,1] ≈ -1200
        @test d[1,2] ≈ 400
        @test d[2,1] ≈ 400
        @test d[2,2] ≈ -800
        bnd = Boundary(Interval(-1,1),noBL,NeumannBC)
        latdom = LatticeDomain(bnd,lat)
        sim = Simulation(10,dom,latdom,fdom)
        h = @test_nowarn Helmholtz(sim; m=2)
        d = @test_nowarn h(0)
        @test d[1,1] ≈ -400
        @test d[1,2] ≈ 400
        @test d[2,1] ≈ 400
        @test d[2,2] ≈ -800
        bnd = Boundary(Interval(-1,1),noBL,MatchedBC{1}(in=[1]),MatchedBC{2}(in=[1]))
        latdom = LatticeDomain(bnd,lat)
        sim = Simulation(10,dom,latdom,fdom)
        h = @test_nowarn Helmholtz(sim; m=2)
        d = @test_nowarn h(20)
    end
    @testset "2D" begin
        cavity = Circle(.4,0,0)
        pump = Square(.3,0,0)
        bnd = Boundary(Rectangle(1,.9,0,0),PML)
        dom = NondispersiveDomain(cavity,2)
        fdom = DispersiveDomain(pump,TwoLevelSystem(10,.1,2))
        lat = Cartesian(.025,.035)
        latdom = LatticeDomain(bnd,lat)
        sim = Simulation(10,dom,latdom,fdom)
        h = @test_nowarn Helmholtz(sim; m=1)
    end
end


text:
module RulesFlowMarginalsTest

using Test
using ReactiveMP
using Random
using Distributions
using LinearAlgebra

import ReactiveMP: @test_marginalrules

@testset "marginalrules:Flow" begin

    @testset ":in (m_out::NormalDistributionsFamily, m_in::NormalDistributionsFamily) (Linearization)" begin

        params = [1.0, 2.0, 3.0]
        model = FlowModel( 2, (AdditiveCouplingLayer(PlanarFlow(); permute=false), ) )
        meta  = FlowMeta(compile(model, params))

        @test_marginalrules [ with_float_conversions = false ] Flow(:in) [
            (
                input = (
                    m_out = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    m_in  = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    m_in  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    m_in = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    m_in = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),
            (
                input = (
                    m_out = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.75002245135493, -2.0000174620747515], [2.5000066522408155 0.5000033261093448; 0.5000033261093448 1.0])
            ),

        ]

    end

    @testset ":in (m_out::NormalDistributionsFamily, m_in::NormalDistributionsFamily) (Unscented)" begin

        params = [1.0, 2.0, 3.0]
        model = FlowModel( 2, (AdditiveCouplingLayer(PlanarFlow(); permute=false), ) )
        meta  = FlowMeta(compile(model, params), Unscented(2))

        @test_marginalrules [ with_float_conversions = false, atol=1e-9 ] Flow(:in) [
            (
                input = (
                    m_out = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    m_in  = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    m_in  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out  = MvNormalMeanPrecision([-5.0, -2.5], diagm([1.0, 1/2.0])),
                    m_in = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    m_in = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),
            (
                input = (
                    m_out = MvNormalWeightedMeanPrecision([-5.0, -1.25], diagm([1.0, 1/2.0])),
                    m_in  = MvNormalMeanCovariance([-5.0, -2.5], diagm([1.0, 2.0])),
                    meta  = meta
                ),
                output = MvNormalWeightedMeanPrecision([-10.750029103730993, -2.0000241143779007], [2.5000066521881763 0.5000033260385859; 0.5000033260385859 0.9999999999115444])
            ),

        ]

    end

end

end


text:
@static if VERSION < v"1.7-"
    ismutabletype(::Type{T}) where T = T.mutable
end


text:
using LinearAlgebra
using DelimitedFiles

function create_basis_mats(dim)

    basis_mats = Any[]

    # Construct diagonal matrices B_ii
    for i = 1:dim
        B =  zeros(Complex{Float64},dim,dim)
        B[i,i] = 1.0
        push!(basis_mats, B)
    end

    # Construct all other mats B_ij
    for i=1:dim
        for j=1:dim
            if i==j
                continue
            end
            B = zeros(Complex{Float64},dim,dim)
            B[i,i] = 1.0/2.0
            B[j,j] = 1.0/2.0

            upp = 1.0/2.0
            low = 1.0/2.0
            if j>i  # imaginary basis mat
                upp =  im * upp
                low = -im * low
            end
            B[i,j] = upp
            B[j,i] = low

            push!(basis_mats, B)
        end
    end

    return basis_mats
end

function sum_basis_mats(dim)
    all = create_basis_mats(dim)

    S =  zeros(Complex{Float64},dim,dim)

    for i = 1:length(all)
        S += all[i]
    end

    S /= dim*dim

    return S
end

function write_mat_to_file(S, filename)

    # Vectorize real and imaginary part and concatenate
    dim = size(S)[1] * size(S)[1]
    Re = reshape(real(S), dim, 1)
    Im = reshape(imag(S), dim, 1)
    x = [Re;Im]
    for i = 1:dim*2
        println(x[i])
        if abs(x[i]) < 1e-12
            x[i] = 0.0
        end
    end

    # Write x to file
    open(filename, "w") do io
        writedlm(io, x)
    end
end

# N=2: for all possible z-coefficients, tests if the linear combination of basis mats is positive definite
function test_N2_coeffs(coeffs_filename)

    N2_coeffs = readdlm(coeffs_filename)
    B = create_basis_mats(2)

    test_passed = true
    for i = 1:size(N2_coeffs)[1]
        eig_i = eigvals(transpose(N2_coeffs[i,:])*B)
        if minimum(eigs) < 0
           test_passed=false
        end
    end

    if test_passed
        print("Yay! Test passed: All linear combinations are positive definite\n")
    else
        print("Mehhhh... some eigenvalues are negative!")
    end
end

max_stars_repo_name: LLNL/scaling-umbrella

max_issues_repo_name: LLNL/scaling-umbrella

max_forks_repo_name: LLNL/scaling-umbrella


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

abstract AbstractEnsemble
abstract DataEnsemble <: AbstractEnsemble
abstract NormalizationEnsemble <: AbstractEnsemble

abstract LatteTask

"""
A task that calls `func` with `args...`
"""
type JuliaTask <: LatteTask
    func :: Function
    args :: Tuple
end

"""
A task that updates parameter `param_id`
"""
type UpdateTask <: LatteTask
    param_id   :: UInt64
end

"""
Used in neuron definitions to mark a field to be unique per batch item
"""
type Batch{T}
    init :: T
end

"""
A parameter in a `Net` (learned during training)

** Fields **

- `name`          -- the name of the parameter
- `gradient_name` -- the name of the gradient (should be ∇`name`)
- `hist_name`     -- the name of the history buffer (should be `name`hist)
- learning_rate   -- local learning rate for the parameter
- regu_coef       -- local regularization coefficient
- clip_gradients  -- NOT IMPLEMENTED, gradient clipping parameter
- value           -- buffer containing the value of the parameter
- gradient        -- buffer containing the gradient of the parameter
- hist            -- buffer containing the history of the parameter
- request         -- request id, used for MPI data parallelism
"""
type Param
    name           :: Symbol
    gradient_name  :: Symbol
    hist_name      :: Symbol
    learning_rate  :: Float32
    regu_coef      :: Float32
    clip_gradients :: Float32

    value    :: Array
    gradient :: Array
    hist     :: Array
    request  :: Cint

    Param(ensemble_name::Symbol, name::Symbol,
          learning_rate::Float32, regu_coef::Float32) =
              new(symbol(ensemble_name, name),
                  symbol(ensemble_name,:∇,name),
                  symbol(ensemble_name, name, :hist), learning_rate, regu_coef, -1.0f0)
end
export Param

"""
A container for tasks for multiple `Phase`s

** Fields **

- tasks -- a dictionary containing a `Vector` of tasks for each `Phase`
"""
type TaskSet
    tasks :: Dict{Phase, Vector{LatteTask}}
    TaskSet() = new(Dict{Phase, Vector{LatteTask}}(Train => [], Test => []))
end

#= function TaskSet(tasks::Dict{Phase, Set}) =#
#=     @assert haskey(tasks, Train) && haskey(tasks, Test) =#
#= end =#

function Base.getindex(task_set::TaskSet, phase::Phase)
    return task_set.tasks[phase]
end

function Base.append!(first::TaskSet, second::TaskSet)
    append!(first[Train], second[Train])
    append!(first[Test], second[Test])
end

type ArgSet
    args :: Dict{Phase, Set}
    ArgSet() = new(Dict{Phase, Set}(Train => Set(), Test => Set()))
end

function ArgSet(args::Dict{Phase, Set})
    @assert haskey(args, Train) && haskey(args, Test)
end

function Base.getindex(arg_set::ArgSet, phase::Phase)
    arg_set.args[phase]
end

function Base.union!(first::ArgSet, second::ArgSet)
    union!(first[Train], second[Train])
    union!(first[Test], second[Test])
end

"""
** Fields **

- `ensembles`          -- a vector containing each `Ensemble` in the network.
- `ensembles_map`      -- a mapping between ensemble names in the network to the
                          instance stored in `ensembles`
- `buffers`            -- the internal buffers used by the network
- `curr_buffer_set`    -- the current buffer set (used for double buffering)
- `forward_tasks`      -- a set of tasks for performing forward propogation of the network
- `backward_tasks`     -- a set of tasks for performing back propogation of the network
- `update_tasks`       -- a set of tasks for performing parameter updates
- `params`             -- a vector of `Param` instances corresponding to network parameters
- `run_where`          -- DEPRECATED
- `signal`             -- DEPRECATED
- `batch_size`         -- the batch size of the network
                          TODO: support different batch sizes for train/test
- `train_epoch`        -- number of epochs completed for training
- `test_epoch`         -- number of epochs completed for testing
- `curr_time_step`     -- the current time step (for RNNs)
- `time_steps`         -- total number of time steps to unroll (for RNNs)
- `num_subgroups`      -- number of partitions in network (for model parallelism)
- `ensemble_send_list` -- a mapping between ensembles and a list of subgroups to send
                          values to, used internally when synthesizing code for
                          model parallelism
"""
type Net
    ensembles          :: Vector{AbstractEnsemble}
    ensembles_map      :: Dict{Symbol, AbstractEnsemble}

    buffers            :: Tuple{Vector{Dict{Symbol, Array}},
                                Vector{Dict{Symbol, Array}}}
    curr_buffer_set    :: Int

    forward_tasks      :: TaskSet
    backward_tasks     :: TaskSet
    update_tasks       :: Vector{LatteTask}

    params             :: Vector{Param}

    run_where          :: Int
    signal             :: Array{Cint, 1}

    batch_size         :: Int

    train_epoch        :: Int
    test_epoch         :: Int

    curr_time_step     :: Int
    time_steps         :: Int

    num_subgroups      :: Int
    ensemble_send_list :: Dict{Symbol, Vector{Tuple{Int, Int}}}

    Net(batch_size, time_steps=1, num_subgroups=1) = new(
                [],                                                  # ensembles
                Dict{Symbol, AbstractEnsemble}(),                    # ensembles_map
                tuple([Dict{Symbol,Array}() for _ in 1:time_steps],
                      [Dict{Symbol,Array}() for _ in 1:time_steps]), # buffers
                1,                                                   # curr_buffer_set
                TaskSet(),                                           # forward_tasks
                TaskSet(),                                           # backward_tasks
                [],                                                  # update_tasks
                Param[],                                             # params
                -1,                                                  # run_where
                Array(Cint, 1),                                      # signal
                batch_size,                                          # batch_size
                1,                                                   # train_epoch
                1,                                                   # test_epoch
                1,                                                   # curr_time_step
                time_steps,                                          # time_steps
                num_subgroups,                                       # num_subgroups
                Dict{Symbol, Vector{Int}}())                         # ensemble_send_list
end

"""
Main `Net` constructor that should be used.
** Params **
- `batch_size`    -- batch size of the network
- `time_steps`    -- number of time steps to unroll the network (for RNNs)
- `num_subgroups` -- number of subgroups in the network (for model parallelism)
"""
function Net(batch_size::Int; time_steps=1, num_subgroups=1)
    net = Net(batch_size, time_steps, num_subgroups)
    @latte_mpi initialize_communicators(net)
    net
end

"""
Returns the batch size of the network
"""
batch_size(net::Net) = net.batch_size

"""
A connection between two ensembles.

** Fields **

- source        -- the source `Ensemble`
- mapping       -- a mapping function to a range of neurons in `source`
- shape         -- shape of the connected neurons returned by `mapping`
- size          -- length of the connected neurons returned by `mapping`
- copy          -- whether the connection requires input values to be copied
- is_dim_fixed  -- vector of booleans that are true if the connection is fixed along a dimension
- is_one_to_one -- whether the connection is one to one
- padding       -- amount of padding used for the connection
- recurrent     -- whether the connection is recurrent
"""
type Connection
    source        :: AbstractEnsemble
    mapping       :: Function
    shape         :: Tuple
    size          :: Int
    copy          :: Bool
    is_dim_fixed  :: Vector{Bool}
    is_one_to_one :: Bool
    padding       :: Int
    recurrent     :: Bool
end

"""
An ensemble

** Fields **

- name         -- name of the ensemble
- neurons      -- an array of neurons of type `T`
- connections  -- a list of `Ensemble`s connected to this ensemble
- batch_fields -- a vector of `Batch` fields for `T` (used internally)
- arg_dim_info --
- params       -- a vector of `Param`s associated with the ensemble
- phase        -- phase(s) in which this ensemble is active
- net_subgroup -- the net subgroup the ensemble is a member of (use for model parallelism)
"""
type Ensemble{T <: Neuron, N} <: AbstractEnsemble
    name         :: Symbol
    neurons      :: Array{T,N}
    connections  :: Vector{Connection}
    batch_fields :: Vector{Symbol}
    arg_dim_info :: Dict{Symbol, Vector{Bool}}
    params       :: Vector
    phase        :: Phase
    net_subgroup :: Cint
end

abstract JuliaEnsemble <: AbstractEnsemble

type ConcatNeuron <: Neuron
    value :: Float32
    ∇     :: Float32
end

type ConcatEnsemble <: AbstractEnsemble
    name :: Symbol
    neurons :: Array{ConcatNeuron}
    connections :: Vector{Connection}
    phase :: Phase
    inputs :: Vector{AbstractEnsemble}
    inner_size :: Int
    net_subgroup :: Int
end

type ReshapeNeuron <: Neuron
    value :: Float32
    ∇     :: Float32
end

type ReshapeEnsemble <: AbstractEnsemble
    name :: Symbol
    neurons :: Array{ReshapeNeuron}
    connections :: Vector{Connection}
    phase :: Phase
    net_subgroup :: Cint
    input :: AbstractEnsemble
    shape :: Tuple
end


text:
# Fast method to compute the derivative of `f` knowing both input and output
derivative(output::Number, f::F, input::Number) where {F} = only(only(derivatives_given_output(output, f, input)))
derivative(output::Number, ::typeof(identity), input::T) where {T<:Number} = one(T)

linear!(y, z, W::AbstractMatrix) = mul!(y, W', z)
linear!(y, z, W::AbstractArray) = conv!(y, z, W, DenseConvDims(z, W))
∇weights!(W::AbstractMatrix, z, ȳ₀) = mul!(W, z, ȳ₀')
∇weights!(W::AbstractArray, z, ȳ₀) = ∇conv_filter!(W, z, ȳ₀, DenseConvDims(z, W))

function flip(W::AbstractArray{T, N}) where {T, N}
    dims = ntuple(identity, N - 2)
    Ŵ = permutedims(W, (dims..., N, N - 1))
    for dim in dims # CUDA does not yet support multiple dims at once, see https://github.com/JuliaGPU/CUDA.jl/issues/1126
        reverse!(Ŵ, dims=dim)
    end
    return Ŵ
end

function clean!(W, project, output_shapes)
    N = length(output_shapes)
    for i0 in 1:N, i1 in i0:N
        W[project(output_shapes[i1], output_shapes[i0])...] .= 0
    end
    return W
end

# `project` must be guaranteed to work for `project(W, input(shape), shape)`
# where `shape` is either in `shapes` or is the whole `y`.
# It must also work for `project(W, shape1, shape2)` with `shape1, shape2 ∈ shapes`.
struct Filtration{N, I, P}
    shapes::Vector{NTuple{N, UnitRange{Int}}}
    input::I
    project::P
end

clean!(W, f::Filtration) = clean!(W, f.project, f.shapes)
clean(W, f::Filtration) = clean!(copy(W), f.project, f.shapes)

const Tensor = AbstractArray{<:Number}
istensor(y) = y isa Tensor

# If `y` is an array of numbers, then `y`, else `z`
macro iftensor(x, y, z=NoTangent())
    return Expr(:if, Expr(:call, :istensor, esc(x)), esc(y), esc(z))
end

apply!(σ, z, y, shape) = view(z, shape...) .+= σ.(view(y, shape...))
apply!(σ::Tensor, z, y, shape) = view(z, shape...) .+= view(σ, shape...) .* view(y, shape...)

# consider `y`, `z` as inputs and solve machine equation in place
function solve!(y, z, W, σ, f::Filtration)
    yshapes, zshapes = f.shapes, map(f.input, f.shapes)
    Wshapes = map(f.project, zshapes, yshapes)
    slices = map(Slices, (y, z, W), (yshapes, zshapes, Wshapes))
    for (i, yshape) in enumerate(yshapes)
        yslice, zslice, Wslice = getindex.(slices, i)
        # linear step (first iteration is trivial)
        add_within!(y => yshape, linear!(yslice, zslice, Wslice) => axes(yslice))
        # nonlinear step
        apply!(σ, z, y, map(intersect, yshape, axes(y)))
    end
    return
end

function init(y₀, z₀, W::AbstractArray{T, N}, forward) where {T, N}
    y = expandchannels(y₀, axes(W, N)) # `y₀` could have fewer channels
    z = isnothing(z₀) ? zero(y) : copy(z₀)
    W̃ = clean(W, forward)
    return (; y, z, W̃)
end

function solve(y₀::Tensor, z₀::Union{Tensor, Nothing}, W, σ, (forward, _))
    y, z, W̃ = init(y₀, z₀, W, forward)
    solve!(y, z, W̃, σ, forward)
    return (; y, z)
end

# simpler method that takes input before nonlinearity and returns output after nonlinearity
solve(y₀::Tensor, W, σ, (forward, backward)) = solve(y₀, nothing, W, σ, (forward, backward)).z

function ChainRulesCore.rrule(::typeof(solve), y₀::Tensor, z₀::Union{Tensor, Nothing}, W, σ, (forward, backward))
    y, z, W̃ = init(y₀, z₀, W, forward)
    solve!(y, z, W̃, σ, forward)
    function solve_pullback(ΔΩ)
        Dσ = @iftensor σ σ derivative.(z .- something(z₀, false), σ, y)
        W′ = flip(W̃)
        Ω̄ = unthunk(ΔΩ)
        ȳ = @iftensor Ω̄.y copy(Ω̄.y) zero(y) # autodiff may pass `AbstractZero`
        z̄ = @iftensor Ω̄.z copy(Ω̄.z) zero(z) # autodiff may pass `AbstractZero`
        solve!(z̄, ȳ, W′, Dσ, backward)
        ȳ₀, z̄₀ = selectchannels(ȳ, axes(y₀)[end-1]), @iftensor z₀ z̄ # Return cotangent comparable to input
        W̄ = ∇weights!(W′, slice(z, forward.input(axes(y))), ȳ) # Reuse `W′` as a buffer
        clean!(W̄, forward) # Compute Σᵢ (πᵢ - πᵢ₋₁) W̄ πᵢ₋₁
        σ̄ = @iftensor σ y .* z̄ # return `NoTangent()` if `σ` is an activation function
        return NoTangent(), ȳ₀, z̄₀, W̄, σ̄, NoTangent()
    end
    return (; y, z), solve_pullback
end