text:
include("orbital_data.jl")
orb_2p_radial = LinearInterpolation(rs_DFT, orb_2p_radial_DFT)

nPts = size(coulomb_precompute)[1]
Rmax = 30;
Rs = range(0, Rmax, length = nPts)
τs = range(0, π / 2, length = nPts)
interp_call = LinearInterpolation((Rs, τs), coulomb_precompute)

"""
    Ψ_pz(r::Float64, θ::Float64)

2pz orbital in spherical coordinates.

The radial part is extracted from DFT, while the angular portion is the
spherical harmonic for 2pz (n=2, l=1, m=0) is cos(θ). The wave function does not
depend on ϕ, as expected for a pz orbital.

# Arguments
* `r`: radial distance from the center of the orbital in Bohr radii
* `θ`: polar angle

Note that the function returns 0 for r > 40 since the radial component at
r = 40 is ≈4.5×10⁻¹¹.
"""
@inline function Ψ_pz(r::Float64, θ::Float64)
    if r == 0.0 || r > 40
        return 0.0
    else
        # Note that the DFT values for the radial portion are r R(r)
        # so we divide by an extra factor of r when computing the WF
        return (orb_2p_radial(r) * cos(θ) / r * √(3 / (4 * π)))
    end
end

"""
    coulomb_potential_pz(R::Float64, τ::Float64)

Coulomb potential generated by a 2pz orbital at distance R and polar angle τ
from the center of the orbital.
"""
function coulomb_potential_pz(R::Float64, τ::Float64)
    # The cubature uses spherical coordinates with x[1] = r, x[2] = θ, x[3] = ϕ
    res = hcubature(
        r ->
            Ψ_pz(r[1], r[2]) .^ 2 * r[1]^2 * sin(r[2]) ./ (√(
                r[1]^2 + R^2 -
                2 *
                r[1] *
                R *
                (cos(τ) * cos(r[2]) + cos(r[3]) * sin(τ) * sin(r[2])),
            )),
        [0, 0, 0],
        [40, π, 2 * π],
        reltol = 1e-4,
    )
end

"""
    coulomb_potential_pz_interp(R::Float64, τ::Float64)

Interpolated version of [`coulomb_potential_pz`](@ref).

Since calculating [`coulomb_potential_pz`](@ref) involves computing a 3D
integral, multiple calculations can be costly. Therefore, this function uses
a precomputed grid to approximate [`coulomb_potential_pz`](@ref) for scenarios
when the direct calculation becomes not feasible.
"""
function coulomb_potential_pz_interp(R::Float64, τ::Float64)
    # Because of the pz shape, the potential is symmetric across the xy plane.
    # Hence, if the polar angle τ > π / 2,  we use τ → π - τ
    # (since the precalculation was performed for 0 < τ < π / 2)
    if τ > π / 2
        τ = π - τ
    end
    if R > Rmax
        return 1 / R
    else
        return interp_call(R, τ)
    end
end

@inline function coulomb_energy_pz_pz_Integrand(
    r::Vector{Float64},
    ρ::Tuple{Float64,Float64,Float64},
)
    # ρ is a CARTESIAN vector between the two pz orbitals
    # r is a SPHERICAL vector pointing from the center of one of the orbitals to
    # a point in space. it is the integration variable
    # dist is a CARTESIAN vector pointing from the center of the second pz to
    # the same point
    dist =
        [ρ[1], ρ[2], ρ[3]] - [
            r[1] * sin(r[2]) * cos(r[3]),
            r[1] * sin(r[2]) * sin(r[3]),
            r[1] * cos(r[2]),
        ]
    res =
        Ψ_pz(r[1], r[2]) .^ 2 *
        r[1]^2 *
        sin(r[2]) *
        coulomb_potential_pz_interp(norm(dist), acos((dist[3]) / norm(dist)))[1]
end

"""
    coulomb_energy_pz_pz(ρ::Tuple{Float64, Float64, Float64})

Coulomb interaction between two graphene orbitals.

# Arguments
* `ρ`: a 3-tuple separation vector (x,y,z) between the orbital centers
"""
function coulomb_energy_pz_pz(ρ::Tuple{Float64,Float64,Float64})
    hcubature(
        r -> coulomb_energy_pz_pz_Integrand(r, ρ),
        [0, 0, 0],
        [40, π, 2 * π],
        reltol = 1e-5,
    )
end


text:
# 추상 타입 Runnable 선언
abstract type Runnable end

# Runnable의 하위 타입 BoomRun 정의
struct BoomRun <: Runnable
    run::Function

    function BoomRun()
        new(function() println("BOOM!!!") end)
    end
end

function main()
    runner::Runnable = BoomRun()
    runner.run()
end

main()

ext: jl

max_stars_repo_name: myyrakle/Julia-Example

max_issues_repo_name: myyrakle/Julia-Example

max_forks_repo_name: myyrakle/Julia-Example


text:
using SafeTestsets

@safetestset "store" begin include("store.jl") end


text:
module DescartesFEA

using TetGen
using Descartes

function tetgen_mesh(m::AbstractMesh, flags)

end

end # module


text:
import Dates
using Dates: Second, Millisecond

warsaw = first(compile("Europe/Warsaw", tzdata["europe"]))
fixed = FixedTimeZone("Fixed", -7200, 3600)

# ZonedDateTime accessors
zdt = ZonedDateTime(DateTime(2014,6,12,23,59,58,57), fixed)
@test timezone(zdt) === fixed

@test TimeZones.days(zdt) == 735396
@test TimeZones.hour(zdt) == 23
@test TimeZones.minute(zdt) == 59
@test TimeZones.second(zdt) == 58
@test TimeZones.millisecond(zdt) == 57

@test eps(zdt) == Millisecond(1)

# Make sure that Dates accessors work with ZonedDateTime.
@test Dates.year(zdt) == 2014
@test Dates.month(zdt) == 6
@test Dates.week(zdt) == 24
@test Dates.day(zdt) == 12
@test Dates.dayofmonth(zdt) == 12
@test Dates.yearmonth(zdt) == (2014, 6)
@test Dates.monthday(zdt) == (6, 12)
@test Dates.yearmonthday(zdt) == (2014, 6, 12)

# Vectorized accessors
# Note: fill is used to test for size and equality.
n = 10
arr = fill(zdt, n)
@test TimeZones.hour.(arr) == fill(23, n)
@test TimeZones.minute.(arr) == fill(59, n)
@test TimeZones.second.(arr) == fill(58, n)
@test TimeZones.millisecond.(arr) == fill(57, n)

@test Dates.year.(arr) == fill(2014, n)
@test Dates.month.(arr) == fill(6, n)
@test Dates.day.(arr) == fill(12, n)
@test Dates.dayofmonth.(arr) == fill(12, n)
@test Dates.yearmonth.(arr) == fill((2014, 6), n)
@test Dates.monthday.(arr) == fill((6, 12), n)
@test Dates.yearmonthday.(arr) == fill((2014, 6, 12), n)


text:
@testset "Rangefinder" begin
    Random.seed!(0)
    m, n = 30, 20
    l = 5
    tol = 1e-8
    B = randn(m, l) * rand(l, n)

    Q = RP.rangefinder(B, 2l; q=0)
    @test norm(B - Q*Q'*B) <= tol

    Q = RP.rangefinder(B, 2l; q=5)
    @test norm(B - Q*Q'*B) <= tol

end


text:
module EMLHairpinIncr
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...streams.EMLstreams
	import EMLtypes.length
	include("HairpinIncr/Properties_Average.jl")
	include("HairpinIncr/Properties_In_Out.jl")
	include("HairpinIncr/Properties_Wall.jl")
	include("HairpinIncr/Physical_Properties.jl")
	include("HairpinIncr/Details_Main.jl")
	include("HairpinIncr/Hairpin_HeatTransfer.jl")
	include("HairpinIncr/Hairpin_PressureDrop.jl")
	include("HairpinIncr/Main_Hairpin.jl")
	include("HairpinIncr/Results_Hairpin.jl")
	include("HairpinIncr/Summary_Hairpin.jl")
	include("HairpinIncr/HairpinIncr_basic.jl")
	include("HairpinIncr/UpperPipe_basic.jl")
	include("HairpinIncr/LowerPipe_basic.jl")
	include("HairpinIncr/HairpinIncr.jl")
end


text:
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

@testset "Preprocessing of the original LP" begin
  @testset "Computing l2 norm of a matrix over the given dimension." begin
    matrix = sparse([3.0 0.0 -4.0; 4.0 3.0 0.0])
    @test FirstOrderLp.l2_norm(matrix, 1) ≈ [5.0, 3.0, 4.0] atol = 1.0e-10
    @test FirstOrderLp.l2_norm(matrix, 2) ≈ [5.0, 5.0] atol = 1.0e-10
  end

  @testset "remove_empty_rows with empty inequality constraint" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        2.0 0.0
        1.0 0.0
        0.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 0.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.remove_empty_rows(problem)
    test_fields_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],               # variable_lower_bound
        [1.0, 2.0],               # variable_upper_bound
        [1.0, 2.0],               # objective_vector
        0.0,                      # objective_constant
        [
          2.0 0.0
          1.0 0.0
        ],                        # constraint_matrix
        [1.0, 1.0],               # right_hand_side
        1,                        # num_equalities
      ),
    )
  end

  @testset "remove_empty_rows with empty equality constraint" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        0.0 0.0
        1.0 0.0
        1.0 0.0
      ],                            # constraint_matrix
      [0.0, 1.0, 0.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.remove_empty_rows(problem)
    test_fields_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],               # variable_lower_bound
        [1.0, 2.0],               # variable_upper_bound
        [1.0, 2.0],               # objective_vector
        0.0,                      # objective_constant
        [
          1.0 0.0
          1.0 0.0
        ],                        # constraint_matrix
        [1.0, 0.0],               # right_hand_side
        0,                        # num_equalities
      ),
    )
  end

  @testset "remove_empty_rows errors if an empty row in an inequality constraint
            has a positive right-hand side" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 0.0
        1.0 0.0
        0.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 1.0],              # right_hand_side
      1,                            # num_equalities
    )
    @test_throws ErrorException FirstOrderLp.remove_empty_rows(problem)
  end

  @testset "remove_empty_rows errors if an empty row in an equality constraint
            has a non-zero right-hand side" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        0.0 0.0
        1.0 0.0
        0.0 1.0
      ],                            # constraint_matrix
      [1.0, 1.0, 1.0],              # right_hand_side
      1,                            # num_equalities
    )
    @test_throws ErrorException FirstOrderLp.remove_empty_rows(problem)
  end

  @testset "remove_empty_columns with removed variable at lower bound" begin
    problem = FirstOrderLp.linear_programming_problem(
      [-1.0, -1.0],                 # variable_lower_bound
      [2.0, 2.0],                   # variable_upper_bound
      [3.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        0.0 1.0
        0.0 -1.0
      ],                            # constraint_matrix
      [1.0, 1.0],                   # right_hand_side
      0,                            # num_equalities
    )
    FirstOrderLp.remove_empty_columns(problem)
    test_fields_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [-1.0],                       # variable_lower_bound
        [2.0],                        # variable_upper_bound
        [2.0],                        # objective_vector
        -3.0,                         # objective_constant
        [[1.0 -1.0]';],               # constraint_matrix
        [1.0, 1.0],                   # right_hand_side
        0,                            # num_equalities
      ),
    )
  end

  @testset "remove_empty_columns with removed variable at upper bound" begin
    problem = FirstOrderLp.linear_programming_problem(
      [-1.0, -1.0],                 # variable_lower_bound
      [2.0, 2.0],                   # variable_upper_bound
      [-3.0, 2.0],                  # objective_vector
      0.0,                          # objective_constant
      [
        0.0 1.0
        0.0 -1.0
      ],                            # constraint_matrix
      [1.0, 1.0],                   # right_hand_side
      0,                            # num_equalities
    )
    FirstOrderLp.remove_empty_columns(problem)
    test_fields_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [-1.0],                       # variable_lower_bound
        [2.0],                        # variable_upper_bound
        [2.0],                        # objective_vector
        -6.0,                         # objective_constant
        [[1.0 -1.0]';],               # constraint_matrix
        [1.0, 1.0],                   # right_hand_side
        0,                            # num_equalities
      ),
    )
  end

  @testset "recover_original_solution" begin
    solution = [1.0, 1.0, 1.0, 5.0]
    empty_indicies = [1, 4]
    original_size = 5
    original_solution = FirstOrderLp.recover_original_solution(
      solution,
      empty_indicies,
      original_size,
    )
    @test original_solution == [0.0, 1.0, 1.0, 0.0, 1.0]
  end

  @testset "presolve" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0, 1.0],              # variable_lower_bound
      [1.0, 2.0, 2.0],              # variable_upper_bound
      [1.0, 2.0, 0.0],              # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0 0.0
        1.0 -1.0 0.0
        0.0 0.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 0.0],              # right_hand_side
      1,                            # num_equalities
    )
    presolve_info = FirstOrderLp.presolve(problem)
    primal_solution, dual_solution =
      FirstOrderLp.undo_presolve(presolve_info, [1.0, 0.0], [1.0, 1.0])
    @test primal_solution == [1.0, 0.0, 1.0]
    @test dual_solution == [1.0, 1.0, 0.0]
  end

  @testset "Presolve doesn't delete empty columns for QP" begin
    problem = FirstOrderLp.QuadraticProgrammingProblem(
      [0.0, 0.0, 0.0],              # variable_lower_bound
      [1.0, 2.0, 1.0],              # variable_upper_bound
      [
        4.0 2.0 0.0
        2.0 1.0 0.0
        0.0 0.0 1.0
      ],                            # objective_matrix
      [1.0, 2.0, 1.0],              # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0 0.0
        1.0 -1.0 0.0
        1.0 0.0 0.0
      ],                            # constraint matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    presolve_info = FirstOrderLp.presolve(problem)
    @test size(problem.constraint_matrix) == (3, 3)
  end

  @testset "lp_norm_rescaling for LP" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0
        1.0 -1.0
        1.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.l2_norm_rescaling(problem)
    # Columns are rescaled by [3^(-1/4), 2^(-1/4)]
    # Rows are rescaled by [2^(-1/4), 2^(-1/4), 1.0]
    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],                         # variable_lower_bound
        [1.0 * (3)^(1 / 4), 2.0 * (2)^(1 / 4)], # variable_upper_bound
        [1.0 / (3)^(1 / 4), 2.0 / (2)^(1 / 4)], # objective_vector
        0.0,                                # objective_constant
        [
          (2*3)^(-1/4) (2*2)^(-1/4)
          (2*3)^(-1/4) -(2 * 2)^(-1 / 4)
          (3)^(-1/4) 0.0
        ],                                  # constraint_matrix
        [(2)^(-1 / 4), (2)^(-1 / 4), 2.0],      # right_hand_side
        1,                                  # num_equalities
      ),
    )
  end

  @testset "lp_norm_rescaling for LP with empty rows" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0
        1.0 -1.0
        0.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 0.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.l2_norm_rescaling(problem)
    # Columns are rescaled by [2^(-1/4), 2^(-1/4)]
    # Rows are rescaled by [2^(-1/4), 2^(-1/4), 1.0]
    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],                         # variable_lower_bound
        [1.0 * (2)^(1 / 4), 2.0 * (2)^(1 / 4)], # variable_upper_bound
        [1.0 / (2)^(1 / 4), 2.0 / (2)^(1 / 4)], # objective_vector
        0.0,                                # objective_constant
        [
          (2*2)^(-1/4) (2*2)^(-1/4)
          (2*2)^(-1/4) -(2 * 2)^(-1 / 4)
          0.0 0.0
        ],                                  # constraint_matrix
        [(2)^(-1 / 4), (2)^(-1 / 4), 0.0],      # right_hand_side
        1,                                  # num_equalities
      ),
    )
  end

  @testset "lp_norm_rescaling for LP with empty columns" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 0.0
        1.0 0.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.l2_norm_rescaling(problem)
    # Columns are rescaled by [6^(-1/4), 1.0]
    # Rows are rescaled by [1.0, 1.0, sqrt(2)]
    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],                         # variable_lower_bound
        [1.0 * (6)^(1 / 4), 2.0], # variable_upper_bound
        [1.0 / (6)^(1 / 4), 2.0], # objective_vector
        0.0,                                # objective_constant
        [
          (6)^(-1/4) 0.0
          (6)^(-1/4) 0.0
          (6*4)^(-1/4)*2.0 0.0
        ],                                  # constraint_matrix
        [1.0, 1.0, 2.0 / sqrt(2)],      # right_hand_side
        1,                                  # num_equalities
      ),
    )
  end

  @testset "Pock-Chambolle rescaling for LP, alpha = 0" begin
    problem = FirstOrderLp.linear_programming_problem(
      [-1.0, -1.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0
        2.0 -1.0
        1.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    constraint_rescaling, variable_rescaling =
      FirstOrderLp.pock_chambolle_rescaling(problem, 0.0)
    @test constraint_rescaling ≈ [sqrt(2), sqrt(2), sqrt(2)]
    @test variable_rescaling ≈ [sqrt(6), sqrt(2)]
  end

  @testset "Pock-Chambolle rescaling for LP, alpha = 1" begin
    problem = FirstOrderLp.linear_programming_problem(
      [-1.0, -1.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0
        2.0 -1.0
        1.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    constraint_rescaling, variable_rescaling =
      FirstOrderLp.pock_chambolle_rescaling(problem, 1.0)
    @test constraint_rescaling ≈ [sqrt(2), sqrt(3), sqrt(1)]
    @test variable_rescaling ≈ [sqrt(4), sqrt(2)]
  end

  @testset "Pock-Chambolle rescaling for LP, alpha = 2" begin
    problem = FirstOrderLp.linear_programming_problem(
      [-1.0, -1.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0
        2.0 -1.0
        1.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    constraint_rescaling, variable_rescaling =
      FirstOrderLp.pock_chambolle_rescaling(problem, 2.0)
    @test constraint_rescaling ≈ [sqrt(2), sqrt(5), sqrt(1)]
    @test variable_rescaling ≈ [sqrt(3), sqrt(3)]
  end

  @testset "Ruiz-Rescaling for LP" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    original_problem = deepcopy(problem)
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 1)
    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],                           # variable_lower_bound
        [sqrt(2), 2.0 * sqrt(3)],             # variable_upper_bound
        [1.0 / sqrt(2), 2.0 / sqrt(3)],       # objective_vector
        0.0,                                  # objective_constant
        [
          1/sqrt(6) 1.0
          0.5 -sqrt(2)/sqrt(3)
          1.0 0.0
        ],                                    # constraint_matrix
        [1 / sqrt(3), 1 / sqrt(2), sqrt(2)],  # right_hand_side
        1,                                    # num_equalities
      ),
    )
    @test cum_variable_rescaling ≈ [2^0.5, 3^0.5]
    @test cum_constraint_rescaling ≈ [3^0.5, 2^0.5, 2^0.5]
    FirstOrderLp.unscale_problem(
      problem,
      cum_constraint_rescaling,
      cum_variable_rescaling,
    )
    test_fields_approx_equal(problem, original_problem)
  end

  @testset "Ruiz Rescaling for LP with empty row and column" begin
    problem = FirstOrderLp.linear_programming_problem(
      [-1.0, -1.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        2.0 0.0
        0.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0],                   # right_hand_side
      1,                            # num_equalities
    )
    original_problem = deepcopy(problem)
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 1)
    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [-sqrt(2), -1.0],           # variable_lower_bound
        [sqrt(2), 2.0],             # variable_upper_bound
        [1 / sqrt(2), 2.0],         # objective_vector
        0.0,                        # objective_constant
        [
          1.0 0.0
          0.0 0.0
        ],                          # constraint_matrix
        [1 / sqrt(2), 1.0],         # right_hand_side
        1,                          # num_equalities
      ),
    )
    @test cum_variable_rescaling ≈ [sqrt(2), 1.0]
    @test cum_constraint_rescaling ≈ [sqrt(2), 1.0]
    FirstOrderLp.unscale_problem(
      problem,
      cum_constraint_rescaling,
      cum_variable_rescaling,
    )
    test_fields_approx_equal(problem, original_problem)
  end

  @testset "Convergence of Ruiz-Rescaling for LP" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 3.0],              # right_hand_side
      1,                            # num_equalities
    )
    original_problem = deepcopy(problem)
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 30)

    constraint_matrix = problem.constraint_matrix

    inf_norm_of_columns = vec(sqrt.(maximum(abs.(constraint_matrix), dims = 1)))
    @test all(inf_norm_of_columns .≈ 1)
    inf_norm_of_rows = vec(sqrt.(maximum(abs.(constraint_matrix), dims = 2)))
    @test all(inf_norm_of_rows .≈ 1)
    FirstOrderLp.unscale_problem(
      problem,
      cum_constraint_rescaling,
      cum_variable_rescaling,
    )
    test_fields_approx_equal(problem, original_problem)
  end

  @testset "rescale_problem for LP" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 3.0],              # right_hand_side
      1,                            # num_equalities
    )
    scaled_problem = FirstOrderLp.rescale_problem(
      10, # l_inf_ruiz_iterations
      true, # l2_norm_rescaling
      nothing, # pock_chambolle_alpha
      0, # verbosity
      problem,
    )
    # The first argument is updated in place.
    FirstOrderLp.unscale_problem(
      scaled_problem.scaled_qp,
      scaled_problem.constraint_rescaling,
      scaled_problem.variable_rescaling,
    )

    test_fields_approx_equal(
      scaled_problem.scaled_qp,
      scaled_problem.original_qp,
    )
  end

  @testset "Ruiz-Rescaling for QP" begin
    problem = FirstOrderLp.QuadraticProgrammingProblem(
      [-Inf, -2.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [4.0 2.0; 2.0 1.0],           # objective_matrix
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    original_problem = deepcopy(problem)
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 1)
    test_fields_approx_equal(
      problem,
      FirstOrderLp.QuadraticProgrammingProblem(
        [-Inf, -2.0 * sqrt(3)],                    # variable_lower_bound
        [2.0, 2.0 * sqrt(3)],                      # variable_upper_bound
        [1.0 1.0/sqrt(3); 1.0/sqrt(3) 1/3],  # objective_matrix
        [0.5, 2.0 / sqrt(3)],                      # objective_vector
        0.0,                                       # objective_constant
        [
          0.5/sqrt(3) 1.0
          0.5/sqrt(2) -sqrt(2)/sqrt(3)
          1.0/sqrt(2) 0.0
        ],                                         # constraint_matrix
        [1 / sqrt(3), 1 / sqrt(2), sqrt(2)],       # right_hand_side
        1,                                         # num_equalities
      ),
    )
    @test cum_variable_rescaling ≈ [2.0, sqrt(3)]
    @test cum_constraint_rescaling ≈ [sqrt(3), sqrt(2), sqrt(2)]
    FirstOrderLp.unscale_problem(
      problem,
      cum_constraint_rescaling,
      cum_variable_rescaling,
    )
    test_fields_approx_equal(problem, original_problem)
  end

  @testset "Convergence of Ruiz-Rescaling for QP" begin
    problem = FirstOrderLp.QuadraticProgrammingProblem(
      [-1.0, -2.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [4.0 2.0; 2.0 1.0],           # objective_matrix
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    original_problem = deepcopy(problem)
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 30)

    constraint_matrix = problem.constraint_matrix
    objective_matrix = problem.objective_matrix

    inf_norm_of_columns = vec(
      sqrt.(
        max.(
          maximum(abs.(constraint_matrix), dims = 1),
          maximum(abs.(objective_matrix), dims = 1),
        ),
      ),
    )
    @test all(inf_norm_of_columns .≈ 1)
    inf_norm_of_rows = vec(sqrt.(maximum(abs.(constraint_matrix), dims = 2)))
    @test all(inf_norm_of_rows .≈ 1)
    FirstOrderLp.unscale_problem(
      problem,
      cum_constraint_rescaling,
      cum_variable_rescaling,
    )
    test_fields_approx_equal(problem, original_problem)
  end

  @testset "L2 Ruiz-Rescaling single iteration for LP" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 3.0],              # right_hand_side
      1,                            # num_equalities
    )
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 1, 2.0)

    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],                               # variable_lower_bound
        [1 * 6^0.25, 2 * 13^0.25],                # variable_upper_bound
        [1 / 6^0.25, 2 / 13^0.25],                # objective_vector
        0.0,                                      # objective_constant
        [
          1/(6*15)^0.25 3/(13*15)^0.25
          1/(7.5*6)^0.25 -2/(13*7.5)^0.25
          2/(6*6)^0.25 0
        ],                                        # constraint_matrix
        [1 / 15^0.25, 1 / 7.5^0.25, 3 / 6^0.25],  # right_hand_side
        1,                                        # num_equalities
      ),
    )

    @test cum_variable_rescaling ≈ [6^0.25, 13^0.25]
    @test cum_constraint_rescaling ≈ [15^0.25, 7.5^0.25, 6^0.25]
  end
  @testset "Convergence of L2 Ruiz-Rescaling for LP" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 3.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.ruiz_rescaling(problem, 60, 2.0)

    norm_of_columns = FirstOrderLp.l2_norm(problem.constraint_matrix, 1)
    @test norm_of_columns ≈ [1, 1] atol = 1e-5
    norm_of_rows = FirstOrderLp.l2_norm(problem.constraint_matrix, 2)
    @test norm_of_rows ≈ [sqrt(2 / 3), sqrt(2 / 3), sqrt(2 / 3)] atol = 1e-5
  end

  @testset "Convergence of L2 Ruiz-Rescaling for QP" begin
    problem = FirstOrderLp.QuadraticProgrammingProblem(
      [-Inf, -2.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [4.0 2.0; 2.0 1.0],           # objective_matrix
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    cum_constraint_rescaling, cum_variable_rescaling =
      FirstOrderLp.ruiz_rescaling(problem, 1, 2.0)
    test_fields_approx_equal(
      problem,
      FirstOrderLp.QuadraticProgrammingProblem(
        [-Inf, -2 * 18^0.25],                       # variable_lower_bound
        [1 * 26^0.25, 2 * 18^0.25],                 # variable_upper_bound
        [
          4/26^0.5 2/(26*18)^0.25
          2/(26*18)^0.25 1/18^0.5
        ],                                          # objective_matrix
        [1 / 26^0.25, 2 / 18^0.25],                 # objective_vector
        0.0,                                        # objective_constant
        [
          1/(25*26)^0.25 3/(18*25)^0.25
          1/(12.5*26)^0.25 -2/(18*12.5)^0.25
          2/(10*26)^0.25 0
        ],                                          # constraint_matrix
        [1 / 25^0.25, 1 / 12.5^0.25, 2 / 10^0.25],  # right_hand_side
        1,                                          # num_equalities
      ),
    )
    @test cum_variable_rescaling ≈ [26^0.25, 18^0.25]
    @test cum_constraint_rescaling ≈ [25^0.25, 12.5^0.25, 10^0.25]
  end

  @testset "Convergence of L2 Ruiz-Rescaling for QP" begin
    problem = FirstOrderLp.QuadraticProgrammingProblem(
      [-1.0, -2.0],                 # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [4.0 2.0; 2.0 1.0],           # objective_matrix
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 3.0
        1.0 -2.0
        2.0 0.0
      ],                            # constraint_matrix
      [1.0, 1.0, 2.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.ruiz_rescaling(problem, 100, 2.0)

    constraint_matrix = problem.constraint_matrix
    objective_matrix = problem.objective_matrix

    norm_of_columns = vec(
      sqrt.(
        sqrt.(
          FirstOrderLp.l2_norm(constraint_matrix, 1) .^ 2 +
          FirstOrderLp.l2_norm(objective_matrix, 1) .^ 2,
        ),
      ),
    )
    @test norm_of_columns ≈ [1, 1] atol = 1e-5
    norm_of_rows = FirstOrderLp.l2_norm(problem.constraint_matrix, 2)
    @test norm_of_rows ≈ [sqrt(2 / 5), sqrt(2 / 5), sqrt(2 / 5)] atol = 1e-5
  end

  # This instance is simple enough to compute where it converges in
  # closed form. (It actually converges in a single iteration.)
  @testset "L2 Ruiz-Rescaling simple" begin
    problem = FirstOrderLp.linear_programming_problem(
      [0.0, 0.0],                   # variable_lower_bound
      [1.0, 2.0],                   # variable_upper_bound
      [1.0, 2.0],                   # objective_vector
      0.0,                          # objective_constant
      [
        1.0 1.0
        1.0 -1.0
        1.0 1.0
      ],                            # constraint_matrix
      [1.0, 1.0, 3.0],              # right_hand_side
      1,                            # num_equalities
    )
    FirstOrderLp.ruiz_rescaling(problem, 10, 2.0)

    test_fields_approx_equal(
      problem,
      FirstOrderLp.linear_programming_problem(
        [0.0, 0.0],                            # variable_lower_bound
        [1 * 3^0.25, 2 * 3^0.25],              # variable_upper_bound
        [1 / 3^0.25, 2 / 3^0.25],              # objective_vector
        0.0,                                   # objective_constant
        [
          1/sqrt(3) 1/sqrt(3)
          1/sqrt(3) -1/sqrt(3)
          1/sqrt(3) 1/sqrt(3)
        ],                                     # constraint_matrix
        [1 / 3^0.25, 1 / 3^0.25, 3 / 3^0.25],  # right_hand_side
        1,                                     # num_equalities
      ),
    )
  end
end


text:
using LambdaTAME: impTTVnodesym, get_kron_contract_comps

@testset "Contraction" begin

    d = 5

    U = rand(A_TOST.n,d)
    V = rand(B_TOST.n,d)
    X = U*V'

    @testset "Kronecker Product Contraction" begin
        UTOST_y = impTTVnodesym(A_UTOST,B_UTOST,reshape(X,A_TOST.n*B_TOST.n))
        #TOST_y = implicit_contraction(A_TOST,B_TOST,reshape(X,A_TOST.n*B_TOST.n))   #TODO: known problem with code, unused in experiment
        TOST_U, TOST_V = get_kron_contract_comps(A_TOST,B_TOST,U,V)
        TOST_X = TOST_U*TOST_V'
        LR_TOST_y = reshape(TOST_X,A_TOST.n*B_TOST.n)
        UST_U, UST_V = get_kron_contract_comps(A_UST,B_UST,U,V)
        UST_X = UST_U*UST_V'
            # UST based embeddings put factors on one
            # embedding rather than using a sqrt

        @test norm(LR_TOST_y - UTOST_y)/norm(UTOST_y) < round_off_bound(UTOST_y)
        @test norm(TOST_X - UST_X)/norm(UST_X) < round_off_bound(UST_X)

    end

    @testset "TAME contraction methods" begin

        A_Ti,B_Ti = LambdaTAME.setup_tame_data(A_TOST,B_TOST)
        A_Mi,B_Mi = LambdaTAME.setup_tame_data(A_UST,B_UST)

        TOST_Y = impTTVnodesym(A_TOST.n,B_TOST.n,size(A_TOST.indices,1),X,A_Ti,B_Ti)
        UST_Y = impTTVnodesym(A_UST.n,B_UST.n,size(A_UST.indices,1),X,A_Mi,B_Mi)

        @test norm(TOST_Y - UST_Y)/norm(UST_Y) < round_off_bound(UST_Y)
    end

    @testset "Multiple Motif Contraction" begin
        #=
        seed!(54321)
        n= 100
        k = 25
        trials = 1000
        orders = [3,4,5,6]

        x = ones(n,1)
        y = zeros(n,1)
        A = LambdaTAME.random_geometric_graph(n,k)
        tensors = tensors_from_graph(A,orders,trials)

        for i =1:length(orders)
            LambdaTAME.contraction_divide_out!(tensors[i],x,y)
        end

        =#
    end

    @testset "SSHOPM" begin

    end

end

max_stars_repo_name: chuckcol/LambdaTAME

max_issues_repo_name: chuckcol/LambdaTAME

max_forks_repo_name: chuckcol/LambdaTAME


text:
module BPPLib

import Base: show

export BPPData, loadBPP, toCSP
export CSPData, loadCSP, toBPP

const data_path = joinpath(pkgdir(BPPLib), "data")

using ZipFile

abstract type BaseData end

include("BPPData.jl")
include("CSPData.jl")
include("Util.jl")
include("Loader.jl")

end


text:
export Star,
       center,
       basis,
       predicate,
       intersection!,
       dim

"""
    Star{N, VN<:AbstractVector{N}, MN<:AbstractMatrix{N}, PT<:AbstractPolyhedron{N}} <: AbstractPolyhedron{N}

Generalized star set with a polyhedral predicate, i.e.

```math
X = \\{x ∈ \\mathbb{R}^n : x = x₀ + \\sum_{i=1}^m α_i v_i,~~\\textrm{s.t. } P(α) = ⊤ \\},
```
where ``x₀ ∈ \\mathbb{R}^n`` is the center, the ``m`` vectors ``v₁, …, vₘ`` form
the basis of the star set, and the combination factors
``α = (α₁, …, αₘ) ∈ \\mathbb{R}^m`` are the predicates' decision variables,
i.e. ``P : α ∈ \\mathbb{R}^m → \\{⊤, ⊥\\}`` where the polyhedral predicate satisfies
``P(α) = ⊤`` if and only if ``Aα ≤ b`` for some fixed ``A ∈ \\mathbb{R}^{p × m}`` and
``b ∈ \\mathbb{R}^p``.

### Fields

- `c` -- vector that represents the center
- `V` -- matrix where each column corresponds to a basis vector
- `P` -- polyhedral set that represents the predicate

### Notes

The predicate function is implemented as a conjunction of linear constraints, i.e.
a subtype of `AbstractPolyhedron`. By a slight abuse of notation, the *predicate*
is also used to denote the subset of ``\\mathbb{R}^n`` such that ``P(α) = ⊤`` holds.

The ``m`` basis vectors (each one ``n``-dimensional) are stored as the columns
of an ``n × m`` matrix.

We remark that a `Star` is mathematically equivalent to the lazy affine map of
the polyhedral set `P`, with the transformation matrix and translation vector
being `V` and `c` respectively.

### Examples

This example is drawn from Example 1 in [2]. Consider the two-dimensional plane
``\\mathbb{R}^2``. Let

```jldoctest star_constructor
julia> V = [[1.0, 0.0], [0.0, 1.0]];
```
be the basis vectors and take

```jldoctest star_constructor
julia> c = [3.0, 3.0];
```
as the center of the star set. Let the predicate be the infinity-norm ball of radius 1,

```jldoctest star_constructor
julia> P = BallInf(zeros(2), 1.0);
```
Finally, the star set ``X = ⟨c, V, P⟩`` defines the set:

```jldoctest star_constructor
julia> S = Star(c, V, P)
Star{Float64, Vector{Float64}, Matrix{Float64}, BallInf{Float64, Vector{Float64}}}([3.0, 3.0], [1.0 0.0; 0.0 1.0], BallInf{Float64, Vector{Float64}}([0.0, 0.0], 1.0))
```

We can use getter functions for each component field:

```jldoctest star_constructor
julia> center(S)
2-element Vector{Float64}:
 3.0
 3.0

julia> basis(S)
2×2 Matrix{Float64}:
 1.0  0.0
 0.0  1.0

julia> predicate(S)
BallInf{Float64, Vector{Float64}}([0.0, 0.0], 1.0)
```
In this case, we know calculating by hand that the generalized star ``S`` is
defined by the rectangular set

```math
    T = \\{(x, y) ∈ \\mathbb{R}^2 : (2 ≤ x ≤ 4) ∧ (2 ≤ y ≤ 4) \\}
```
It holds that ``T`` and ``S`` are equivalent.

### References

Star sets as defined here were introduced in [1]; see also [2] for a preliminary
definition. For applications in reachability analysis of neural networks, see
[3].

- [1] Duggirala, P. S., and Mahesh V. *Parsimonious, simulation based verification of linear systems.*
      International Conference on Computer Aided Verification. Springer, Cham, 2016.

- [2] Bak S, Duggirala PS. *Simulation-equivalent reachability of large linear systems with inputs.*
      In International Conference on Computer Aided Verification
      2017 Jul 24 (pp. 401-420). Springer, Cham.

- [3] Tran, H. D., Lopez, D. M., Musau, P., Yang, X., Nguyen, L. V., Xiang, W., & Johnson, T. T. (2019, October).
      *Star-based reachability analysis of deep neural networks.*
      In International Symposium on Formal Methods (pp. 670-686). Springer, Cham.
"""
struct Star{N, VN<:AbstractVector{N}, MN<:AbstractMatrix{N}, PT<:AbstractPolyhedron{N}} <: AbstractPolyhedron{N}
    c::VN # center
    V::MN # basis
    P::PT # predicate

    # default constructor with size checks
    function Star(c::VN, V::MN, P::PT) where {N, VN<:AbstractVector{N}, MN<:AbstractMatrix{N}, PT<:AbstractPolyhedron{N}}
        @assert length(c) == size(V, 1) "the center of the basis vectors should be compatible, " *
                                        "but they are of length $(length(c)) and $(size(V, 1)) respectively"

        @assert dim(P) == size(V, 2) "the number of basis vectors should be compatible " *
                                     "with the predicates' dimension, but they are $(size(V, 2)) and $(dim(P)) respectively"

        return new{N, VN, MN, PT}(c, V, P)
    end
end

# constructor from center and list of generators
function Star(c::VN, vlist::AbstractVector{VN}, P::PT) where {N, VN<:AbstractVector{N}, PT<:AbstractPolyhedron{N}}
    V = to_matrix(vlist, length(c))
    return Star(c, V, P)
end

# type used for dispatch
const STAR{N, VN<:AbstractVector{N},
              MN<:AbstractMatrix{N},
              PT<:AbstractPolyhedron{N}} = AffineMap{N, PT, N, MN, VN}

isoperationtype(::Type{<:Star}) = false
isconvextype(::Type{<:Star}) = true

# ============================
# Star set getter functions
# ============================

"""
    center(X::Star)

Return the center of a star.

### Input

- `X` -- star

### Output

The center of the star.
"""
center(X::Star) = X.c

"""
    basis(X::Star)

Return the basis vectors of a star.

### Input

- `X` -- star

### Output

A matrix where each column is a basis vector of the star.
"""
basis(X::Star) = X.V

"""
    predicate(X::Star)

Return the predicate of a star.

### Input

- `X` -- star

### Output

A polyhedral set representing the predicate of the star.
"""
predicate(X::Star) = X.P

# =============================
# LazySets interface functions
# =============================

"""
    dim(X::Star)

Return the dimension of a star.

### Input

- `X` -- star

### Output

The dimension of a star.
"""
function dim(X::Star)
    return length(X.c)
end

"""
    σ(d::AbstractVector, X::Star)

Return the support vector of a star.

### Input

- `d` -- direction
- `X` -- star

### Output

The support vector in the given direction.
"""
function σ(d::AbstractVector, X::Star)
    A = basis(X)
    return A * σ(_At_mul_B(A, d), predicate(X)) + center(X)
end

"""
    ρ(d::AbstractVector, X::Star)

Return the support function of a star.

### Input

- `d` -- direction
- `X` -- star

### Output

The support function in the given direction.
"""
function ρ(d::AbstractVector, X::Star)
    return ρ(_At_mul_B(basis(X), d), predicate(X)) + dot(d, center(X))
end

"""
    an_element(X::Star)

Return some element of an affine map.

### Input

- `X` -- star

### Output

An element of the star. It relies on the `an_element` function of the
wrapped set.
"""
function an_element(X::Star)
    return basis(X) * an_element(predicate(X)) + center(X)
end

"""
    isempty(X::Star)

Return whether an affine map is empty or not.

### Input

- `am` -- affine map

### Output

`true` iff the predicate is empty.
"""
function isempty(X::Star)
    return isempty(predicate(X))
end

"""
    isbounded(am::Star; cond_tol::Number=DEFAULT_COND_TOL)

Determine whether a star is bounded.

### Input

- `X`        -- star
- `cond_tol` -- (optional) tolerance of matrix condition (used to check whether
                the basis matrix is invertible)

### Output

`true` iff the star is bounded.

### Algorithm

See [`isbounded(::AbstractAffineMap)`](@ref).
"""
function isbounded(X::Star; cond_tol::Number=DEFAULT_COND_TOL)
    am = convert(STAR, X)
    return isbounded(am, cond_tol=cond_tol)
end

"""
    ∈(v::AbstractVector, X::Star)

Check whether a given point is contained in a star set.

### Input

- `v`  -- point/vector
- `X` -- star

### Output

`true` iff ``v ∈ X``.

### Algorithm

See [`∈(::AbstractVector, ::AbstractAffineMap)`](@ref).
"""
function ∈(x::AbstractVector, X::Star)
    return basis(X) \ (x - center(X)) ∈ predicate(X)
end

"""
    vertices_list(X::Star; apply_convex_hull::Bool=true)

Return the list of vertices of a star set.

### Input

- `X`                 -- star
- `apply_convex_hull` -- (optional, default: `true`) if `true`, apply the convex
                         hull operation to the list of vertices of the star

### Output

A list of vertices.

### Algorithm

See [`vertices_list(::AbstractAffineMap)`](@ref).
"""
function vertices_list(X::Star; apply_convex_hull::Bool=true)
    am = convert(STAR, X)
    return vertices_list(am, apply_convex_hull=apply_convex_hull)
end

"""
    constraints_list(X::Star)

Return the list of constraints of a star.

### Input

- `X` -- star

### Output

The list of constraints of the star.

### Algorithm

See [`constraints_list(::AbstractAffineMap)`](@ref).
"""
function constraints_list(X::Star)
    am = convert(STAR, X)
    return constraints_list(am)
end

"""
    linear_map(M::AbstractMatrix, X::Star)

Return the linear map of a star.

### Input

- `M` -- matrix
- `X` -- star

### Output

The star obtained by applying `M` to `X`.
"""
function linear_map(M::AbstractMatrix, X::Star)
    c′ = M * X.c
    V′ = M * X.V
    P′ = X.P
    return Star(c′, V′, P′)
end

"""
    affine_map(M::AbstractMatrix, X::Star, v::AbstractVector)

Return the affine map of a star.

### Input

- `M` -- matrix
- `X` -- star
- `v` -- vector

### Output

The star obtained by the affine map with matrix `M` and displacement `v` to `X`.
"""
function affine_map(M::AbstractMatrix, X::Star, v::AbstractVector)
    c′ = M * X.c + v
    V′ = M * X.V
    P′ = X.P
    return Star(c′, V′, P′)
end

"""
    rand(::Type{Star}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing)

Create a random star.

### Input

- `Star`       -- type for dispatch
- `N`              -- (optional, default: `Float64`) numeric type
- `dim`            -- (optional, default: 2) dimension
- `rng`            -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`           -- (optional, default: `nothing`) seed for reseeding

### Output

A random star.

### Algorithm

This functions calls the function to generate a random zonotope,
because every zonotope is a star.
"""
function rand(::Type{Star};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing)
    rng = reseed(rng, seed)
    Z = rand(Zonotope, N=N, dim=dim, rng=rng, seed=seed)
    X = convert(Star, Z)
    return X
end


text:
@testset "Persistence" begin
    @test_throws ErrorException LightGraphs._NI("Not implemented")
    pdict = loadgraphs(joinpath(testdir,"testdata","tutte-pathdigraph.jgz"))
    p1 = pdict["Tutte"]
    p2 = pdict["pathdigraph"]
    g3 = PathGraph(5)

    (f,fio) = mktemp()
    # test :lg
    @test savegraph(f, p1) == 1
    @test savegraph(f, p1; compress=true) == 1
    @test savegraph(f, p1, LGFormat(); compress=true) == 1
    @test savegraph(f, p2; compress=true) == 1
    @test (ne(p2), nv(p2)) == (9, 10)

    g2 = loadgraph(f)
    @test (ne(g2), nv(g2)) == (9, 10)

    @test length(sprint(savegraph, p1, LGFormat())) == 401
    @test length(sprint(savegraph, p2, LGFormat())) == 50
    gs = loadgraph(joinpath(testdir, "testdata", "tutte-pathdigraph.jgz"), "pathdigraph")
    @test gs == p2
    @test_throws ErrorException loadgraph(joinpath(testdir, "testdata", "tutte-pathdigraph.jgz"), "badname")

    d = Dict{String, AbstractGraph}("p1"=>p1, "p2"=>p2)
    @test savegraph(f, d) == 2

end

max_stars_repo_name: AsileBcd/JohnsonCycles

max_issues_repo_name: AsileBcd/JohnsonCycles

max_forks_repo_name: AsileBcd/JohnsonCycles


text:
"""
    XB.jl

This is a Julia port of a MATLAB implementation of batch and incremental
Xie-Beni (XB) Cluster Validity Index.

Authors:
MATLAB implementation: Leonardo Enzo Brito da Silva
Julia port: Sasha Petrenko <sap625@mst.edu>

REFERENCES
[1] X. L. Xie and G. Beni, "A Validity Measure for Fuzzy Clustering," IEEE
Transactions on Pattern Analysis and Machine Intelligence, vol. 13, no. 8,
pp. 841–847, 1991.
[2] M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, and J. Bailey,
"Online Cluster Validity Indices for Streaming Data," ArXiv e-prints, 2018,
arXiv:1801.02937v1 [stat.ML]. [Online].
[3] M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, J. Bailey, "Online
cluster validity indices for performance monitoring of streaming data clustering,"
Int. J. Intell. Syst., pp. 1–23, 2018.
"""

"""
    XB

The stateful information of the Xie-Beni (XB) Cluster Validity Index.

# References
1. X. L. Xie and G. Beni, "A Validity Measure for Fuzzy Clustering," IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 13, no. 8, pp. 841–847, 1991.
2. M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, and J. Bailey, "Online Cluster Validity Indices for Streaming Data," ArXiv e-prints, 2018, arXiv:1801.02937v1 [stat.ML]. [Online].
3. M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, J. Bailey, "Online cluster validity indices for performance monitoring of streaming data clustering," Int. J. Intell. Syst., pp. 1–23, 2018.
"""
mutable struct XB <: CVI
    label_map::LabelMap
    dim::Integer
    n_samples::Integer
    mu_data::RealVector     # dim
    n::IntegerVector        # dim
    v::RealMatrix           # dim x n_clusters
    CP::RealVector          # dim
    SEP::Float
    G::RealMatrix           # dim x n_clusters
    D::RealMatrix           # n_clusters x n_clusters
    WGSS::Float
    n_clusters::Integer
    criterion_value::Float
end # XB <: CVI

"""
    XB()

Default constructor for the Xie-Beni (XB) Cluster Validity Index.
"""
function XB()
    XB(
        LabelMap(),                     # label_map
        0,                              # dim
        0,                              # n_samples
        Array{Float, 1}(undef, 0),      # mu_data
        Array{Integer, 1}(undef, 0),    # n
        Array{Float, 2}(undef, 0, 0),   # v
        Array{Float, 1}(undef, 0),      # CP
        0.0,                            # SEP
        Array{Float, 2}(undef, 0, 0),   # G
        Array{Float, 2}(undef, 0, 0),   # D
        0.0,                            # WGSS
        0,                              # n_clusters
        0.0                             # criterion_value
    )
end # XB()

"""
    setup!(cvi::XB, sample::Vector{T}) where {T<:RealFP}
"""
function setup!(cvi::XB, sample::Vector{T}) where {T<:RealFP}
    # Get the feature dimension
    cvi.dim = length(sample)
    # Initialize the 2-D arrays with the correct feature dimension
    cvi.v = Array{T, 2}(undef, cvi.dim, 0)
    cvi.G = Array{T, 2}(undef, cvi.dim, 0)
end # setup!(cvi::XB, sample::Vector{T}) where {T<:RealFP}

function param_inc!(cvi::XB, sample::RealVector, label::Integer)
    # Get the internal label
    i_label = get_internal_label!(cvi.label_map, label)

    n_samples_new = cvi.n_samples + 1
    if isempty(cvi.mu_data)
        mu_data_new = sample
        setup!(cvi, sample)
    else
        mu_data_new = (1 - 1/n_samples_new) .* cvi.mu_data + (1/n_samples_new) .* sample
    end

    if i_label > cvi.n_clusters
        n_new = 1
        v_new = sample
        CP_new = 0.0
        G_new = zeros(cvi.dim)
        if cvi.n_clusters == 0
            D_new = zeros(1, 1)
        else
            D_new = zeros(cvi.n_clusters + 1, cvi.n_clusters + 1)
            D_new[1:cvi.n_clusters, 1:cvi.n_clusters] = cvi.D
            d_column_new = zeros(cvi.n_clusters + 1)
            # println(d_column_new)
            for jx = 1:cvi.n_clusters
                # sum((v_new - cvi.v[:, jx]).^2, dims=1)
                d_column_new[jx] = sum((v_new - cvi.v[:, jx]).^2)
            end
            D_new[:, i_label] = d_column_new
            D_new[i_label, :] = transpose(d_column_new)
        end
        # Update 1-D parameters with a push
        cvi.n_clusters += 1
        push!(cvi.CP, CP_new)
        push!(cvi.n, n_new)
        # Update 2-D parameters with appending and reassignment
        cvi.v = [cvi.v v_new]
        cvi.G = [cvi.G G_new]
        cvi.D = D_new
    else
        n_new = cvi.n[i_label] + 1
        v_new = (1 - 1/n_new) .* cvi.v[:, i_label] + (1/n_new) .* sample
        delta_v = cvi.v[:, i_label] - v_new
        diff_x_v = sample .- v_new
        CP_new = cvi.CP[i_label] + transpose(diff_x_v)*diff_x_v + cvi.n[i_label]*transpose(delta_v)*delta_v + 2*transpose(delta_v)*cvi.G[:, i_label]
        d_column_new = zeros(cvi.n_clusters)
        for jx = 1:cvi.n_clusters
            if jx == i_label
                continue
            end
            # sum((v_new - cvi.v[:, jx]).^2, dims=1)
            d_column_new[jx] = sum((v_new - cvi.v[:, jx]).^2)
        end
        # Update parameters
        cvi.n[i_label] = n_new
        cvi.v[:, i_label] = v_new
        cvi.CP[i_label] = CP_new
        cvi.D[:, i_label] = d_column_new
        cvi.D[i_label, :] = transpose(d_column_new)
    end
    cvi.n_samples = n_samples_new
    cvi.mu_data = mu_data_new
end # param_inc!(cvi::XB, sample::RealVector, label::Integer)

function param_batch!(cvi::XB, data::RealMatrix, labels::IntegerVector)
    cvi.dim, cvi.n_samples = size(data)
    # Take the average across all samples, but cast to 1-D vector
    cvi.mu_data = mean(data, dims=2)[:]
    # u = findfirst.(isequal.(unique(labels)), [labels])
    u = unique(labels)
    cvi.n_clusters = length(u)
    cvi.n = zeros(Integer, cvi.n_clusters)
    cvi.v = zeros(cvi.dim, cvi.n_clusters)
    cvi.CP = zeros(cvi.n_clusters)
    cvi.D = zeros(cvi.n_clusters, cvi.n_clusters)
    for ix = 1:cvi.n_clusters
        subset = data[:, findall(x->x==u[ix], labels)]
        cvi.n[ix] = size(subset, 2)
        cvi.v[1:cvi.dim, ix] = mean(subset, dims=2)
        diff_x_v = subset - cvi.v[:, ix] * ones(1, cvi.n[ix])
        cvi.CP[ix] = sum(diff_x_v.^2)
    end
    for ix = 1 : (cvi.n_clusters - 1)
        for jx = ix + 1 : cvi.n_clusters
            # cvi.D[jx, ix] = sum((cvi.v[:, ix] - cvi.v[:, jx]).^2, dims=1)
            cvi.D[jx, ix] = sum((cvi.v[:, ix] - cvi.v[:, jx]).^2)
        end
    end
    cvi.D = cvi.D + transpose(cvi.D)
end # param_batch!(cvi::XB, data::RealMatrix, labels::IntegerVector)

function evaluate!(cvi::XB)
    cvi.WGSS = sum(cvi.CP)
    if cvi.n_clusters > 1
        # Assume a symmetric dimension
        dim = size(cvi.D)[1]
        # Get the values from D as the upper triangular offset from the diagonal
        values = [cvi.D[i,j] for i = 1:dim, j=1:dim if j > i]
        # SEP is the minimum of these unique D values
        cvi.SEP = minimum(values)
        # Criterion value is
        cvi.criterion_value = cvi.WGSS/(cvi.n_samples*cvi.SEP)
    else
        cvi.SEP = 0.0
        cvi.criterion_value = 0.0
    end
end # evaluate(cvi::XB)


text:
abstract type UStatistics end

########################################

struct OneSampleUstatistics{N} <: UStatistics
  h
end

struct TwoSampleUstatistics{N, M} <: UStatistics
  h
end

######################################

compute(s::UStatistics, X::AbstractArray) = compute(s, s.h, X)
compute(s::UStatistics, X::AbstractArray, Y::AbstractArray) = compute(s, s.h, X, Y)

function compute(::OneSampleUstatistics{1}, h, X::AbstractVector)
  v = zero(eltype(X))
  n = 0
  for x in X
    n += 1
    v += h(x)
  end

  v / n
end

function compute(::OneSampleUstatistics{1}, h, X::AbstractMatrix)
  n = size(X, 1)
  v = 0.
  @inbounds for i=1:n
    v += h(view(X, i, :))
  end

  v / n
end

function compute(::OneSampleUstatistics{2}, h, X::AbstractVector)
  n = length(X)
  v = 0.
  @inbounds for i1=1:n
    for i2=i1+1:n
      v += h( X[i1], X[i2] )::Float64
    end
  end

  2. * v / (n * (n-1))
end

function compute(::OneSampleUstatistics{2}, h, X::AbstractMatrix)
  n = size(X, 1)
  v = 0.
  @inbounds for i1=1:n
    for i2=i1+1:n
      v += h( view(X, i1, :),
              view(X, i2, :) )
    end
  end

  2. * v / n / (n-1)
end

function compute(::TwoSampleUstatistics{1, 1}, h, X::AbstractVector, Y::AbstractVector)
  nx = length(X)
  ny = length(Y)

  v = 0.
  for i=1:nx
    for j=1:ny
      v += h(X[i], Y[j])
    end
  end
  v / nx / ny
end

function compute(s::TwoSampleUstatistics{1, 1}, X::AbstractMatrix, Y::AbstractMatrix)
  h = s.h
  nx = size(X, 1)
  ny = size(Y, 1)

  v = 0.
  for i=1:nx
    for j=1:ny
      v += h(view(X, i, :), view(Y, j, :))
    end
  end
  v / nx / ny
end

######################################

variance(s::UStatistics, X::AbstractArray) = variance(s, s.h, X)
variance(s::UStatistics, X::AbstractArray, Y::AbstractArray) = variance(s, s.h, X, Y)

function variance(::OneSampleUstatistics{1}, h, X::AbstractVector)
  n = length(X)
  q = map(h, X)
  t = mean(q)
  (sum(abs2, q) - n * t^2) / (n - 1)
end

function variance(::OneSampleUstatistics{1}, h, X::AbstractMatrix)
  n = size(X, 1)
  q = mapslices(h, X; dims = 2)
  t = mean(q)
  (sum(abs2, q) - n * t^2) / (n - 1)
end

function variance(::TwoSampleUstatistics{1, 1}, h, X::AbstractMatrix, Y::AbstractMatrix)
  nx = size(X, 1)
  ny = size(Y, 1)

  q = zeros(nx)
  @inbounds for k=1:nx
    vx = view(X, k, :)
    for j=1:ny
      q[k] += h(vx, view(Y, j, :))
    end
    q[k] /= ny
  end

  r = zeros(ny)
  @inbounds for k=1:ny
    vy = view(Y, k, :)
    for j=1:nx
      r[k] += h(view(X, j, :), vy)
    end
    r[k] /= nx
  end

  t = mean(q)

  σ1 = sum(abs2, q) / (nx - 1) - nx / (nx - 1) * t^2
  σ2 = sum(abs2, r) / (ny - 1) - ny / (ny - 1) * t^2

  return σ1/nx + σ2/ny
end

######################################

### examples

function createVarianceStatistic()
  @inline function h(x, y)
    (x - y)^2 / 2.
  end
  OneSampleUstatistics{2}(h)
end

function createCovarianceStatistic(row::Int, col::Int)
  @inline function h(X, Y)
      @inbounds (X[row] - Y[row]) * (X[col] - Y[col]) / 2.
  end
  OneSampleUstatistics{2}(h)
end

function createSecondMomentStatistic()
  h(X) = X*X
  OneSampleUstatistics{1}(h)
end

function createSecondMomentStatistic(row::Int, col::Int)
  @inbounds h(X) = X[row] * X[col]
  OneSampleUstatistics{1}(h)
end


text:
module LogExpFunctions

using DocStringExtensions: SIGNATURES
using Base: Math.@horner

import ChainRulesCore
import ChangesOfVariables
import InverseFunctions
import IrrationalConstants
import LinearAlgebra

export xlogx, xlogy, xlog1py, xexpx, xexpy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,
    softplus, invsoftplus, log1pmx, logmxp1, logaddexp, logsubexp, logsumexp, softmax,
    softmax!, logcosh

include("basicfuns.jl")
include("logsumexp.jl")
include("chainrules.jl")
include("inverse.jl")
include("with_logabsdet_jacobian.jl")

end # module


text:
using Gadfly, DataArrays, RDatasets

set_default_plot_size(6inch, 3inch)

plot(dataset("Zelig", "macro"), x="Year", y="Country", color="GDP",
     Coord.cartesian(raster=true), Geom.rectbin)


text:
#=
   Port of my the Or-tools/Python Sudoku solver.

   After warmup:
   julia> @time include("sudoku_or_tools.jl")
    num_solutions: 1
    [8, 6, 1, 4, 5, 9, 7, 2, 3]
    [4, 5, 2, 3, 1, 7, 6, 9, 8]
    [7, 9, 3, 6, 8, 2, 5, 1, 4]
    [2, 1, 6, 8, 3, 5, 4, 7, 9]
    [9, 8, 4, 2, 7, 6, 1, 3, 5]
    [3, 7, 5, 1, 9, 4, 8, 6, 2]
    [5, 4, 7, 9, 2, 1, 3, 8, 6]
    [1, 3, 9, 5, 6, 8, 2, 4, 7]
    [6, 2, 8, 7, 4, 3, 9, 5, 1]

    num_solutions:1
    failures:0
    branches:0
    WallTime:291
    0.297074 seconds (571.88 k allocations: 29.771 MiB, 3.09% gc time)

   Via command line:
   $ time julia sudoku_or_tools.jl
    num_solutions: 1
    [8, 6, 1, 4, 5, 9, 7, 2, 3]
    [4, 5, 2, 3, 1, 7, 6, 9, 8]
    [7, 9, 3, 6, 8, 2, 5, 1, 4]
    [2, 1, 6, 8, 3, 5, 4, 7, 9]
    [9, 8, 4, 2, 7, 6, 1, 3, 5]
    [3, 7, 5, 1, 9, 4, 8, 6, 2]
    [5, 4, 7, 9, 2, 1, 3, 8, 6]
    [1, 3, 9, 5, 6, 8, 2, 4, 7]
    [6, 2, 8, 7, 4, 3, 9, 5, 1]

    num_solutions:1
    failures:0
    branches:0
    WallTime:839
    julia sudoku_or_tools.jl  2,63s user 0,78s system 155% cpu 2,193 total

=#
using PyCall

# Resize matrix to "normal" form
function resize_matrix(grid)
    rows = size(grid)
    transpose(reduce(hcat,grid)) # Suggestion by @Mason Protter @ Zulip
end

# Be sure that ortools are installed correctly
# Note that we are using the old CP solver (not the new hot SAT solver)
ot = pyimport("ortools.constraint_solver.pywrapcp")

# Create the solver
solver = ot.Solver("Sudoku")

function sudoku(initial_grid)
    cell_size = 3
    line_size = cell_size^2
    line = 1:line_size
    cell = 1:cell_size

    # Decision variables
    grid = [solver.IntVar(1, line_size, "grid[$i,$j]") for i in line, j in line]

    # AllDifferent on rows.
    for i in line
        solver.Add(solver.AllDifferent([grid[i, j] for j in line]))
    end

    # AllDifferent on columns.
    for j in line
        solver.Add(solver.AllDifferent([grid[i, j] for i in line]))
    end

    # AllDifferent on cells.
    for i in cell, j in cell
        solver.Add(solver.AllDifferent(
                        vec([grid[(i-1) * cell_size + di, (j-1) * cell_size + dj] for di in cell, dj in cell] )
                    ))
    end

    # Initial values.
    for i in line
        for j in line
            if initial_grid[i,j] > 0
                solver.Add(grid[i, j] == initial_grid[i,j])
            end
        end
    end

    # Regroup all variables into a list.
    all_vars = [grid[i, j] for i in line for j in line]

    # Create search phases.
    vars_phase = solver.Phase(all_vars,
                                solver.INT_VAR_SIMPLE,
                                solver.INT_VALUE_SIMPLE)

    solution = solver.Assignment()
    solution.Add(all_vars)
    # Note: This give just one solution, even if there are many others
    collector = solver.FirstSolutionCollector(solution)

    # And solve.
    solver.Solve(vars_phase, [collector])

    num_solutions = collector.SolutionCount()
    println("num_solutions: $num_solutions")
    if num_solutions == 1
        for i in line
            println([collector.Value(0, grid[i, j]) for j in line])
        end
    end
    println("\nnum_solutions:", num_solutions)
    println("failures:", solver.Failures())
    println("branches:", solver.Branches())
    println("WallTime:", solver.WallTime())

end

initial_grid = resize_matrix([[0, 6, 0, 0, 5, 0, 0, 2, 0],
    [0, 0, 0, 3, 0, 0, 0, 9, 0],
    [7, 0, 0, 6, 0, 0, 0, 1, 0],
    [0, 0, 6, 0, 3, 0, 4, 0, 0],
    [0, 0, 4, 0, 7, 0, 1, 0, 0],
    [0, 0, 5, 0, 9, 0, 8, 0, 0],
    [0, 4, 0, 0, 0, 1, 0, 0, 6],
    [0, 3, 0, 0, 0, 8, 0, 0, 0],
    [0, 2, 0, 0, 4, 0, 0, 5, 0]])

sudoku(initial_grid)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
# # Problems and the system

#md # ```@meta
#md # CurrentModule = ImmersedLayers
#md # ```

#=
In specific problems that we wish to solve with immersed layers, there may
be other data and operators that we would like to cache. We do this with
an *extra cache*, which the user can define, along with a problem type associated
with this cache. The basic cache and the extra cache are generated and associated
together in a *system*.

There are a few basic ingredients to do this:
* **Create a problem type**, using the macro `@ilmproblem`. This mostly just serves as a means of dispatching correctly, but also holds the grid and bodies so they can be passed along when the caches are constructed.
* **Create an extra cache type**, making it a subtype of `AbstractExtraILMCache`. This can hold pretty much anything you want it to.
* **Extend the function `prob_cache(prob,base_cache)`** to serve as a constructor for your extra cache, when your problem type is passed in.

Optionally, you can also extend the function `solve` in order to perform the steps
of your algorithm. However, generically, you can just use pass in the system structure,
which holds the basic ILM cache and your extra cache, into any function.
=#

#=
## Example of problem and system use
We will demonstrate the use of problems and systems with the example
given in [A Dirichlet Poisson problem](@ref). Here, we will
assemble the various additional data structures and operators used to
solve this problem into an extra cache. We will also create a problem
type called `DirichletPoissonProblem`, which we make a subtype of
`AbstractScalarILMProblem`.
=#

using ImmersedLayers
using Plots
using UnPack

#=
### Create your problem type
For this job, we have a macro
=#
@ilmproblem DirichletPoisson scalar

#=
This generates a type `DirichletPoissonProblem` for dispatch, and an associated
constructor that we'll use later.
=#

#=
### Create your extra cache
Here, we'd like this extra cache to hold the Schur complement and the
filtering matrices, as well as some cache variables.
=#
struct DirichletPoissonCache{SMT,CMT,ST,FT} <: ImmersedLayers.AbstractExtraILMCache
   S :: SMT
   C :: CMT
   fb :: ST
   fstar :: FT
end

#=
### Extend the `prob_cache` function
We need this to construct our extra cache
=#
function ImmersedLayers.prob_cache(prob::DirichletPoissonProblem,base_cache::BasicILMCache)
    S = create_RTLinvR(base_cache)
    C = create_surface_filter(base_cache)
    fb = zeros_surface(base_cache)
    fstar = zeros_grid(base_cache)
    DirichletPoissonCache(S,C,fb,fstar)
end

#=
### Extend the `solve` function
Here, we actually do the work of the algorithm, making use of all of the
operators and data structures that we have cached for efficiency.
The example below takes in some surface Dirichlet data `fbplus`,
and returns the solutions `f` and `s` (filtered).
=#
function ImmersedLayers.solve(fbplus,prob::DirichletPoissonProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack S, C, fb, fstar = extra_cache

    f = zeros_grid(base_cache)
    s = zeros_surface(base_cache)

    surface_divergence!(fstar,fbplus,base_cache)
    fb .= 0.5*fbplus

    inverse_laplacian!(fstar,base_cache)

    interpolate!(s,fstar,base_cache)
    s .= fb - s
    s .= -(S\s);

    regularize!(f,s,base_cache)
    inverse_laplacian!(f,base_cache)
    f .+= fstar;

    return f, C^6*s
end

#=
### Set up the grid, shape, and cache
We do this just as we did in [Immersed layer caches](@ref), but
now we don't create a cache, since it will be done internally.
=#
Δx = 0.01
Lx = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx)
RadC = 1.0
Δs = 1.4*cellsize(g)
body = Circle(RadC,Δs)

#=
### Do the work
We do this in three steps:
- Create the problem instance. This is where we use the constructor generated by `@ilmproblem`
- Call `__init` to create the caches, assembled into a system
- Call `solve` to solve the problem.
Also, note that pretty much any function that accepts `base_cache`
as an argument also accepts `sys`.
=#
prob = DirichletPoissonProblem(g,body,scaling=GridScaling)
sys = ImmersedLayers.__init(prob)

pts = points(sys)
f, s = solve(pts.u,prob,sys)

plot(f,sys)
#-
plot(s)

#=
In the previous example, we supplied the surface data as an argument directly to
the `solve` function. However, sometimes we need to supply surface data with a little
more complexity (as in time-varying cases). For this purpose, we can create
functions that supply the boundary data when called, e.g.,
=#
get_fbplus(base_cache) = points(base_cache).u
get_fbminus(base_cache) = zeros_surface(base_cache)

#=
We can pass these along to the problem specification by various means.
It's at the user's discretion how to do it, and how to make use of them
in the solution. Here's an example, using a `Dict`
=#
bcdict = Dict("fbplus"=>get_fbplus,"fbminus"=>get_fbminus)

#=
We can also pass along physical parameters and forcing functions.
Let's use the latter feature in order to set some point sources at
a few locations. For this, we create a function that will take in
the base cache and the right-hand side data vector, and populate
this vector with the point sources. We use the `Regularize` function
in the `CartesianGrids.jl` package to immerse the point sources
into the grid. In the example shown here, two points: one at (-1.0,1.5)
and the other at (1.0,1.5), with respective strengths -1 and 1,
are placed.
=#
function rhs_func!(rhs,base_cache)
    X = VectorData([-1.0,1.0],[1.5,1.5])
    str = ScalarData([-1.0,1.0])
    reg = Regularize(X,cellsize(base_cache.g),I0=origin(base_cache.g),ddftype=CartesianGrids.M4prime)
    return reg(rhs,str)
end

#=
We redefine the `solve` function and use these functions in place of
the original argument:
=#
function ImmersedLayers.solve(prob::DirichletPoissonProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack bc, f_funcs, gdata_cache = base_cache
    @unpack S, C, fb, fstar = extra_cache

    f = zeros_grid(base_cache)
    s = zeros_surface(base_cache)

    ## Get the boundary data on each side of the interface
    fbplus = bc["fbplus"](base_cache)
    fbminus = bc["fbminus"](base_cache)

    ## Evaluate the forcing field
    f_funcs(gdata_cache,base_cache)

    ## Evaluate the right-hand side of Poisson equation
    surface_divergence!(fstar,fbplus-fbminus,base_cache)
    fstar .+= gdata_cache

    fb .= 0.5*(fbplus+fbminus)

    inverse_laplacian!(fstar,base_cache)

    interpolate!(s,fstar,base_cache)
    s .= fb - s
    s .= -(S\s);

    regularize!(f,s,base_cache)
    inverse_laplacian!(f,base_cache)
    f .+= fstar;

    return f, C^6*s
end

#=
Now we specify the problem, create the system, and solve it, as before,
but now supplying the boundary condition functions with the `bc` keyword:
=#
prob = DirichletPoissonProblem(g,body,scaling=GridScaling,bc=bcdict,f_funcs=rhs_func!)
sys = ImmersedLayers.__init(prob)
f, s = solve(prob,sys)
plot(f,sys)

#=
Same solution, of course. But suppose we wish to change the
the boundary conditions or source points? We can do it easily without regenerating the
cache and system, simply by redefining our bc and forcing functions, e.g.
to create an internal solution, with surface data equal to the $y$ coordinate,
=#
get_fbplus(base_cache) = zeros_surface(base_cache)
get_fbminus(base_cache) = points(base_cache).v
function rhs_func!(rhs,base_cache)
    X = VectorData([-0.2,0.2],[0.0,0.0])
    str = ScalarData([-1.0,1.0])
    reg = Regularize(X,cellsize(base_cache.g),I0=origin(base_cache.g),ddftype=CartesianGrids.M4prime)
    return reg(rhs,str)
end
#=
We can solve immediately without having to reconstruct the system, so
it's very fast.
=#
@time f, s = solve(prob,sys)
plot(f,sys)

#md # ## Problem types and functions
#md #
#md # ```@docs
#md # AbstractScalarILMProblem
#md # AbstractVectorILMProblem
#md # @ilmproblem
#md # BasicScalarILMProblem
#md # BasicVectorILMProblem
#md # prob_cache
#md # ```

#md # ## System types and functions
#md #
#md # ```@docs
#md # ILMSystem
#md # __init
#md # ```


text:
#=
function inv_2(b::T) where {T<:AbstractFloat}
    hi = inv(b)
    v = hi * b
    w = fma(hi, b, -v)
    lo = (one(T) - v - w) / b
    return hi, lo
end
=#

function inv_2(b::T) where {T<:AbstractFloat}
     hi = inv(b)
     lo = fma(hi, b, -1.0)
     lo = -lo / b
     return hi, lo
end

@inline inv_(b::T) where {T<:AbstractFloat} = inv_2(b)

#=
function div_(a::T, b::T) where {T<:AbstractFloat}
    hi = a / b
    v = hi * b
    w = fma(hi, b, -v)
    lo = (a - v - w) / b
    return hi, lo
end
=#
# !?! `y` must be negated to get the right result

function dvi_2(a::T, b::T) where {T<:AbstractFloat}
     hi = a / b
     lo = -(fma(hi, b, -a) / b)
     return hi, lo
end

@inline dvi_(a::T, b::T) where {T<:AbstractFloat} = dvi_2(a,b)

function root2_2(a::T) where {T<:AbstractFloat}
    hi = sqrt(a)
    lo = fma(-hi, hi, a) / (hi + hi)
    return hi, lo
end

@inline root2_(x::T) where {T<:AbstractFloat} = root2_2(x)

#=
"Concerning the division, the elementary rounding error is
generally not a floating point number, so it cannot be computed
exactly. Hence we cannot expect to obtain an error
free transformation for the xdivision. ...
This means that the computed approximation is as good as
we can expect in the working precision."
-- http://perso.ens-lyon.fr/nicolas.louvet/LaLo05.pdf

While the sqrt algorithm is not strictly an errorfree transformation,
it is known to be reliable and is recommended for general use.
"Augmented precision square roots, 2-D norms and
   discussion on correctly reounding xsqrt(x^2 + y^2)"
by Nicolas Brisebarre, Mioara Joldes, Erik Martin-Dorel,
   Hean-Michel Muller, Peter Kornerup
=#


text:
#
# Returns a 2D array of distributions. The ij-th element is for the
# j-th element of the i-th particle from the previous population
#
function generate_kernels(
        population::AbstractArray{F,2},
        priors::AbstractArray{D,1},
        ) where {
        F<:Real,
        D<:ContinuousUnivariateDistribution,
        }
    n_particles = size(population, 1)
    n_params = size(population, 2)

    if n_particles > 1
        stds = reshape(std(population, dims=1), n_params)
    else
        stds = 1e-3 * ones(n_params) # If there is only one particle we cannot compute the sd - use a small value instead?
    end

    lowers = minimum.(priors)
    uppers = maximum.(priors)

    kernels = Matrix{ContinuousUnivariateDistribution}(undef, n_particles, n_params)
    for j in 1:n_params
        means = population[:, j]
        kernels[:, j] = TruncatedNormal.(means, stds[j]*sqrt(2.0), lowers[j], uppers[j])
    end

    return kernels
end

# For batch of parameters
function generate_parameters(
        batch_size::Int64,
        priors::AbstractArray{D1,1},
        old_weights::StatsBase.Weights,
        kernels::AbstractArray{D2,2},
        ) where {
        D1<:ContinuousUnivariateDistribution,
        D2<:ContinuousUnivariateDistribution
        }

    n_params = length(priors)

    # ADD DimensionMismatch THROWS SO @inbounds CAN BE USED?

    # the kernels must be centered around the old particles
    # and truncated to the priors.

    particles = StatsBase.sample(axes(kernels, 1), old_weights, batch_size)
    perturbed_parameters = rand.(kernels[particles,:])

    # gives a batch_size x n_params matrix of prior pdfs in perturbed parameters
    pdfs = pdf.(reshape(priors, 1, n_params), perturbed_parameters)
    numerators = prod(pdfs, dims=2) # multiply across rows, to get a column vector of products of size batch_size

    denominators = zeros(length(numerators), 1)
    for k in eachindex(denominators)
        denominator_pdfs = pdf.(kernels, reshape(perturbed_parameters[k, :], 1, n_params))
        denominator_summands = prod(denominator_pdfs, dims=2)
        denominators[k] = sum(old_weights .* denominator_summands)
    end

    weight_value = reshape(numerators ./ denominators, batch_size)

    return perturbed_parameters, weight_value
end

# For single parameter only - return parameter as 1D array and weight as float
function generate_parameters(
    priors::AbstractArray{D1,1},
    old_weights::StatsBase.Weights,
    kernels::AbstractArray{D2,2},
    ) where {
    D1, D2<:ContinuousUnivariateDistribution
    }

    perturbed_parameters, weight_value = generate_parameters(1, priors, old_weights, kernels)
    return perturbed_parameters, weight_value[1]
end

function generate_parameters_no_weights(
        n_batch_size::Int64,
        old_parameters::AbstractArray{F,2},
        old_weights::StatsBase.Weights,
        kernels::AbstractArray{D2,2}
        ) where {
        D2<:ContinuousUnivariateDistribution,
        F<:AbstractFloat,
        }
    particles = StatsBase.sample(axes(old_parameters, 1), old_weights, n_batch_size)
    return rand.(kernels[particles,:])
end

function normalise(
        weights::StatsBase.AbstractWeights;
        tosum = 1.0,
        )
    WeightType = typeof(weights)
    weights = WeightType(weights.values .* (tosum / sum(weights.values)), tosum)

    return weights
end

function normalise(weight_values::AbstractArray{F, 1}, tosum = 1.0) where F<:AbstractFloat
    return StatsBase.Weights(weight_values .* (sum(weight_values) / tosum), tosum)
end

#
# Initialise a simulated ABC-SMC run
#
function initialiseABCSMC(input::SimulatedABCSMCInput;
        write_progress = true,
        progress_every = 1000,
        )
    # the first run is an ABC rejection simulation
    rejection_input = SimulatedABCRejectionInput(input)

    rejection_output = ABCrejection(rejection_input;
                                    write_progress = write_progress,
                                    progress_every = progress_every)

    tracker =  SimulatedABCSMCTracker(input.n_params,
                             [rejection_output.n_accepted],
                             [rejection_output.n_tries],
                             [rejection_output.threshold],
                             rejection_output.n_accepted > 0,
                             [rejection_output.population],
                             [rejection_output.distances],
                             [rejection_output.weights],
                             input.priors,
                             input.distance_simulation_input,
                             input.max_iter
                             )

    return tracker
end

#
# Initialise an emulated ABC-SMC run
#
function initialiseABCSMC(input::EmulatedABCSMCInput{CUD, ER, EPS};
        write_progress = true,
        progress_every = 1000) where
        {CUD<:ContinuousUnivariateDistribution,
        ER<:AbstractEmulatorRetraining,
        EPS<:AbstractEmulatedParticleSelection}

    # the first run is an ABC rejection simulation
    rejection_input = EmulatedABCRejectionInput(input.n_params,
                                        input.n_particles,
                                        input.threshold_schedule[1],
                                        input.priors,
                                        input.batch_size,
                                        input.max_iter,
                                        input.emulator_training_input,
                                        input.selection)
    rejection_output = ABCrejection(rejection_input;
                                    write_progress = write_progress)

    tracker = EmulatedABCSMCTracker{CUD, typeof(rejection_output.emulator), ER, EPS}(input.n_params,
                             [rejection_output.n_accepted],
                             [rejection_output.n_tries],
                             [rejection_output.threshold],
                             rejection_output.n_accepted > 0,
                             [rejection_output.population],
                             [rejection_output.distances],
                             [rejection_output.weights],
                             input.priors,
                             input.emulator_training_input,
                             input.emulator_retraining,
                             input.selection,
                             input.batch_size,
                             input.max_iter,
                             [rejection_output.emulator]
                             )

    return tracker
end

#
# Iterate a simulated ABC-SMC
#
function iterateABCSMC!(tracker::SimulatedABCSMCTracker,
        threshold::AbstractFloat,
        n_toaccept::Int;
        write_progress = true,
        progress_every = 1000)

    if write_progress
        @info "GpABC SMC simulation ϵ = $threshold"
    end
    if threshold > tracker.threshold_schedule[end]
        @warn "current threshold less strict than previous one."
    end

    # initialise
    kernels = generate_kernels(tracker.population[end], tracker.priors)
    population = zeros(n_toaccept, tracker.n_params)
    distances = zeros(n_toaccept)
    weight_values = zeros(n_toaccept)
    distance = 0.0
    n_tries = 0
    n_accepted = 0

    # simulate
    while n_accepted < n_toaccept && n_tries < tracker.max_iter
        parameters, weight_value = generate_parameters(tracker.priors, tracker.weights[end], kernels)
        # run simulation for a single particle
        try
            distance = simulate_distance(parameters, tracker.distance_simulation_input)
        catch e
            if isa(e, DimensionMismatch)
                # This prevents the whole code from failing if there is a problem
                # solving the differential equation(s). The exception is thrown by the
                # distance function
                @warn "The summarised simulated data does not have the same size as the summarised reference data. If this is not happening at every iteration it may be due to the behaviour of DifferentialEquations::solve - please check for related warnings. Continuing to the next iteration."
                n_tries += 1
                continue
            else
                throw(e)
            end
        end

        n_tries += 1

        # Handle result
        if distance[1] <= threshold
            n_accepted += 1
            population[n_accepted,:] = parameters
            distances[n_accepted] = distance[1]
            weight_values[n_accepted] = weight_value
        end

        if write_progress && (n_tries % progress_every == 0)
            @info "GpABC SMC simulation accepted $(n_accepted)/$(n_tries) particles."
        end
    end

    if n_accepted == 0
        @warn "Simulation reached maximum $(tracker.max_iter) iterations without selecting any particles"
    elseif n_accepted < n_toaccept
        population = population[1:n_accepted, :]
        weight_values = weight_values[1:n_accepted]
        distances = distances[1:n_accepted]
        @warn "Simulation reached maximum $(tracker.max_iter) iterations before finding $(n_toaccept) particles - will return $n_accepted"
    end

    update_smctracker!(tracker, n_accepted, n_tries, threshold,
                        population, distances, weight_values)
end

#
# Iterate a emulated ABC-SMC
#
function iterateABCSMC!(tracker::EmulatedABCSMCTracker{CUD, ET, ER, EPS},
        threshold::AbstractFloat,
        n_toaccept::Int;
        write_progress = true,
        progress_every = 1000
        ) where {CUD<:ContinuousUnivariateDistribution, ET, ER<:AbstractEmulatorRetraining, EPS<:AbstractEmulatedParticleSelection}
    if write_progress
        @info "GpABC SMC emulation ϵ = $threshold."
    end

    if threshold > tracker.threshold_schedule[end]
        @warn "Current threshold less strict than previous one."
    end

    # initialise
    iter_no = 0
    n_accepted = 0
    n_tries = 0
    old_population = tracker.population[end]
    old_weights = tracker.weights[end]
    population = zeros(n_toaccept, tracker.n_params)
    all_distances=zeros(n_toaccept)
    all_weight_values=zeros(n_toaccept)
    kernels = generate_kernels(tracker.population[end], tracker.priors)
    particle_sampling_function(batch_size) = generate_parameters_no_weights(batch_size,
        old_population,
        old_weights,
        kernels)

    emulator = abc_retrain_emulator(tracker.emulators[end], particle_sampling_function, threshold,
        tracker.emulator_training_input, tracker.emulator_retraining)

    # emulate
    while n_accepted < n_toaccept && iter_no < tracker.max_iter
        parameters, weight_values = generate_parameters(tracker.batch_size,
                                                 tracker.priors,
                                                 old_weights,
                                                 kernels)
        n_tries += size(parameters, 1)
        distances, accepted_indices = abc_select_emulated_particles(emulator, parameters, threshold, tracker.selection)
        n_include = length(accepted_indices)
        if n_accepted + n_include > n_toaccept
            n_include = n_toaccept - n_accepted
            accepted_indices = accepted_indices[1:n_include]
            distances = distances[1:n_include]
        end
        store_slice = n_accepted + 1 : n_accepted + n_include
        all_distances[store_slice] = distances
        population[store_slice,:] = parameters[accepted_indices, :]
        all_weight_values[store_slice] = weight_values[accepted_indices]
        n_accepted += n_include

        if write_progress
            @info "GpABC SMC emulation accepted $(n_accepted)/$(n_tries) particles."
        end

        iter_no += 1
    end

    if n_accepted == 0
        @warn "Emulation reached maximum $(tracker.max_iter) iterations without selecting any particles"
    elseif n_accepted < n_toaccept
        population = population[1:n_accepted, :]
        all_weight_values = all_weight_values[1:n_accepted]
        all_distances = all_distances[1:n_accepted]
        @warn "Emulation reached maximum $(tracker.max_iter) iterations before finding $(n_toaccept) particles - will return $n_accepted"
    end

    update_smctracker!(tracker, n_accepted, n_tries, threshold,
                        population, all_distances, all_weight_values, emulator)
end

function buildAbcSmcOutput(tracker::EmulatedABCSMCTracker)
    EmulatedABCSMCOutput(length(tracker.priors),
                        tracker.n_accepted,
                        tracker.n_tries,
                        tracker.threshold_schedule,
                        tracker.population,
                        tracker.distances,
                        tracker.weights,
                        tracker.emulators)
end

function buildAbcSmcOutput(tracker::SimulatedABCSMCTracker)
    SimulatedABCSMCOutput(length(tracker.priors),
                        tracker.n_accepted,
                        tracker.n_tries,
                        tracker.threshold_schedule,
                        tracker.population,
                        tracker.distances,
                        tracker.weights)
end

function ABCSMC(
        input::T;
        write_progress = true,
        progress_every = 1000,
        ) where {T<:ABCSMCInput}

    tracker = initialiseABCSMC(input; write_progress = write_progress, progress_every = progress_every)

    if tracker.can_continue
        for i in 2:length(input.threshold_schedule)
            threshold = input.threshold_schedule[i]
            iterateABCSMC!(tracker, threshold, input.n_particles;
                write_progress = write_progress,
                progress_every = progress_every)
            if !tracker.can_continue
                break
            end
        end
    else
        @warn "No particles selected at initial rejection ABC step of simulated SMC ABC - terminating algorithm"
    end

    return buildAbcSmcOutput(tracker)
end

# not exported
function initialise_abcsmc_iteration(tracker::ABCSMCTracker, n_toaccept::Int)

    kernels = generate_kernels(tracker.population[end], tracker.priors)
    population = zeros(n_toaccept, tracker.n_params)
    distances = zeros(n_toaccept)
    weight_values = zeros(n_toaccept)

    return kernels, population, weight_values, distances
end

# not exported
function update_smctracker!(tracker::SimulatedABCSMCTracker, n_accepted::Int,
    n_tries::Int, threshold::AbstractFloat, population::AbstractArray{Float64,2},
    distances::AbstractArray{Float64,1}, weight_values::AbstractArray{Float64,1})

    if n_accepted > 0
        tracker.can_continue = true
        push!(tracker.n_accepted, n_accepted)
        push!(tracker.n_tries, n_tries)
        push!(tracker.threshold_schedule, threshold)
        push!(tracker.population, population)
        push!(tracker.distances, distances)
        # Do not want to normalise weights now if doing model selection - will do at
        # end of population at model selection level
        push!(tracker.weights, normalise(weight_values))
    else
        tracker.can_continue = false
    end
end

# not exported
function update_smctracker!(tracker::EmulatedABCSMCTracker, n_accepted::Int,
    n_tries::Int, threshold::AbstractFloat, population::AbstractArray{Float64,2},
    distances::AbstractArray{Float64,1}, weight_values::AbstractArray{Float64,1},
    emulator::GPModel)

    if n_accepted > 0
        tracker.can_continue = true
        push!(tracker.n_accepted, n_accepted)
        push!(tracker.n_tries, n_tries)
        push!(tracker.threshold_schedule, threshold)
        push!(tracker.population, population)
        push!(tracker.distances, distances)
        # Do not want to normalise weights now if doing model selection - will do at
        # end of population at model selection level
        push!(tracker.weights, normalise(weight_values))
        push!(tracker.emulators, emulator)
    else
        tracker.can_continue = false
    end
end


text:
using Test
using KernelFunctions
using Distances
using Random

@testset "KernelFunctions" begin
    include("test_kernelmatrix.jl")
    include("test_approximations.jl")
    include("test_constructors.jl")
    # include("test_AD.jl")
    include("test_transform.jl")
    include("test_distances.jl")
    include("test_kernels.jl")
    include("test_generic.jl")
    include("test_adjoints.jl")
    include("test_custom.jl")
    #include("types.jl")
end


text:
using Test
using LinearAlgebra
using DFTK
import DFTK: total_local_potential
include("testcases.jl")

# TODO Once we have converged SCF densities in a file it would be better to instead / also
#      test the energies of these densities and compare them directly to the reference
#      energies obtained in the data files

@testset "Using BZ symmetry yields identical density" begin
    function get_bands(testcase, kcoords, ksymops, symops, atoms; Ecut=5, tol=1e-8, n_rounds=1)
        kwargs = ()
        n_bands = div(testcase.n_electrons, 2)
        if testcase.Tsmear !== nothing
            kwargs = (temperature=testcase.Tsmear, smearing=DFTK.Smearing.FermiDirac())
            n_bands = div(testcase.n_electrons, 2) + 4
        end

        model = model_DFT(testcase.lattice, atoms, :lda_xc_teter93; kwargs...)
        basis = PlaneWaveBasis(model, Ecut, kcoords, ksymops, symops)
        ham = Hamiltonian(basis; ρ=guess_density(basis, atoms))

        res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol=tol)
        occ, εF = DFTK.find_occupation(basis, res.λ)
        ρnew = compute_density(basis, res.X, occ)

        for it in 1:n_rounds
            ham = Hamiltonian(basis; ρ=ρnew)
            res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol=tol, guess=res.X)

            occ, εF = DFTK.find_occupation(basis, res.λ)
            ρnew = compute_density(basis, res.X, occ)
        end

        ham, res.X, res.λ, ρnew, occ
    end

    function test_orthonormality(basis, ψ; tol=1e-8)
        n_k = length(ψ)
        n_states = size(ψ[1], 2)
        n_fft = prod(basis.fft_size)

        for (ik, kpt) in enumerate(basis.kpoints)
            # Fourier-transform the wave functions to real space
            ψk = ψ[ik]
            ψk_real = cat((DFTK.G_to_r(basis, kpt, ψik) for ψik in eachcol(ψk))..., dims=4)

            T = real(eltype(ψk_real))
            ψk_real_mat = reshape(ψk_real, n_fft, n_states)
            ψk_real_overlap = adjoint(ψk_real_mat) * ψk_real_mat
            nondiag = ψk_real_overlap - I * (n_fft / basis.model.unit_cell_volume)

            @test maximum(abs.(nondiag)) < tol
        end
    end

    function test_full_vs_irreducible(testcase, kgrid_size; Ecut=5, tol=1e-8, n_ignore=0,
                                      kshift=[0, 0, 0], eigenvectors=true)
        spec = ElementPsp(testcase.atnum, psp=load_psp(testcase.psp))
        atoms = [spec => testcase.positions]

        kfull, sym_full, symops = bzmesh_uniform(kgrid_size, kshift=kshift)
        res = get_bands(testcase, kfull, sym_full, symops, atoms; Ecut=Ecut, tol=tol)
        ham_full, ψ_full, eigenvalues_full, ρ_full, occ_full = res
        test_orthonormality(ham_full.basis, ψ_full, tol=tol)

        kcoords, ksymops, symops = bzmesh_ir_wedge(kgrid_size, DFTK.symmetry_operations(testcase.lattice, atoms), kshift=kshift)
        res = get_bands(testcase, kcoords, ksymops, symops, atoms; Ecut=Ecut, tol=tol)
        ham_ir, ψ_ir, eigenvalues_ir, ρ_ir, occ_ir = res
        test_orthonormality(ham_ir.basis, ψ_ir, tol=tol)
        @test ham_full.basis.fft_size == ham_ir.basis.fft_size

        # Test density is the same in both schemes, and symmetric wrt the basis symops
        @test maximum(abs.(ρ_ir.fourier - ρ_full.fourier)) < 10tol
        @test maximum(abs.(ρ_ir.real - ρ_full.real)) < 10tol
        @test maximum(abs, DFTK.symmetrize(ρ_ir; symops=symops).fourier - ρ_ir.fourier) < tol

        # Test local potential is the same in both schemes
        @test maximum(abs, total_local_potential(ham_ir) - total_local_potential(ham_full)) < tol

        # Test equivalent k-Points have the same orbital energies
        for (ik, k) in enumerate(kcoords)
            for (S, τ) in ksymops[ik]
                ikfull = findfirst(1:length(kfull)) do idx
                    all(isinteger, kfull[idx] - S * k)
                end
                @test ikfull !== nothing

                # Orbital energies should agree (by symmetry)
                # The largest few are ignored, because the results between the
                # full kpoints and the reduced kpoints are sometimes a little unstable
                # (since we actually use different guesses)
                range = 1:length(eigenvalues_full[ikfull]) - n_ignore
                @test eigenvalues_full[ikfull][range] ≈ eigenvalues_ir[ik][range] atol=tol
            end
        end

        if eigenvectors
            # Test applying the symmetry transformation to the irreducible k-Points
            # yields an eigenfunction of the Hamiltonian
            # Also check that the accumulated partial densities are equal
            # to the returned density.
            ρsum = zero(ρ_ir.fourier)
            n_ρ = 0
            for (ik, k) in enumerate(kcoords)
                Hk_ir = ham_ir.blocks[ik]
                for (S, τ) in ksymops[ik]
                    Skpoint, ψSk = DFTK.apply_ksymop((S, τ), ham_ir.basis, Hk_ir.kpoint, ψ_ir[ik])

                    ikfull = findfirst(1:length(kfull)) do idx
                        all(isinteger, round.(kfull[idx] - Skpoint.coordinate, digits=10))
                    end
                    @test !isnothing(ikfull)
                    Hk_full = ham_full.blocks[ikfull]

                    range = 1:length(eigenvalues_ir[ik]) - n_ignore
                    for iband in range
                        ψnSk = ψSk[:, iband]
                        residual = norm(Hk_full * ψnSk - eigenvalues_ir[ik][iband] * ψnSk)
                        @test residual < 10tol
                    end  # iband

                    n_ρ += 1
                    ρsum .+= DFTK.compute_partial_density(ham_ir.basis, Skpoint, ψSk, occ_ir[ik])
                end  # (S, τ)
            end  # k

            @test n_ρ == length(kfull)
            @test maximum(abs, ρsum / n_ρ - ρ_full.fourier) < 10tol
        end # eigenvectors
    end

    test_full_vs_irreducible(silicon, [3, 2, 3], Ecut=5, tol=1e-10)

    test_full_vs_irreducible(silicon, [3, 3, 3], Ecut=5, tol=1e-6)
    test_full_vs_irreducible(silicon, [3, 3, 3], Ecut=5, tol=1e-6, kshift=[1/2, 1/2, 1/2])
    test_full_vs_irreducible(silicon, [3, 3, 3], Ecut=5, tol=1e-6, kshift=[1/2, 0, 1/2])
    test_full_vs_irreducible(silicon, [3, 3, 3], Ecut=5, tol=1e-6, kshift=[0, 1/2, 0])

    test_full_vs_irreducible(silicon, [2, 3, 4], Ecut=5, tol=1e-6)
    test_full_vs_irreducible(magnesium, [2, 3, 4], Ecut=5, tol=1e-6, n_ignore=1)
    test_full_vs_irreducible(aluminium, [1, 3, 5], Ecut=3, tol=1e-5, n_ignore=3,
                             eigenvectors=false)

    # That's pretty expensive:
    # test_full_vs_irreducible([4, 4, 4], Ecut=5, tol=1e-6)
end