text:
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#
# Problem set 5 solutions
# Written by Tyler Ransom
# Commented by Giuseppe Grasso
# Recording available in Notability
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#

using Random
using LinearAlgebra
using Statistics
using Optim
using DataFrames
using DataFramesMeta
using CSV
using HTTP
using GLM
cd("/Users/peppegrass/Documents/GitHub/fall-2020/ProblemSets/PS5-ddc/") # GG: sets working directory
pwd() ## GG: prints working directory
readdir() # GG: equivalent to -ls- to see elements of working directory

# read in function to create state transitions for dynamic model
include("create_grids.jl")

function allwrap()

    #=
    GG: In this problem set, we will explore a simpliﬁed version of the Rust (1987, Econometrica) bus engine
    replacement model. Let’s start by reading in the data.
    =#

# GG: I - PRELIMINARIES: Loading data and reshaping them

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 1: Reshaping the data
    # Reshape the data into “long” panel format, calling your long dataset df_long.
    # I have included code on how to do this in the PS5starter.jl ﬁle that accompanies this problem set.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # load in the data
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdataBeta0.csv"
    df = CSV.read(HTTP.get(url).body)

    # create bus id variable
    df = @transform(df, bus_id = 1:size(df,1))

    #---------------------------------------------------
    # reshape from wide to long (must do this twice be-
    # cause DataFrames.stack() requires doing it one
    # variable at a time)
    #---------------------------------------------------
    # first reshape the decision variable
    dfy = @select(df, :bus_id,:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20,:RouteUsage,:Branded)
    dfy_long = DataFrames.stack(dfy, Not([:bus_id,:RouteUsage,:Branded]))
    rename!(dfy_long, :value => :Y)
    dfy_long = @transform(dfy_long, time = kron(collect([1:20]...),ones(size(df,1))))
    select!(dfy_long, Not(:variable))

    # next reshape the odometer variable
    dfx = @select(df, :bus_id,:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20)
    dfx_long = DataFrames.stack(dfx, Not([:bus_id]))
    rename!(dfx_long, :value => :Odometer)
    dfx_long = @transform(dfx_long, time = kron(collect([1:20]...),ones(size(df,1))))
    select!(dfx_long, Not(:variable))

    # join reshaped df's back together
    df_long = leftjoin(dfy_long, dfx_long, on = [:bus_id,:time])
    sort!(df_long,[:bus_id,:time])

# GG: II - STATIC ESTIMATION

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 2: Estimate a STATIC version of the model
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    #=
    GG: The model we would like to estimate is Harold Zurcher’s decision to run buses in his ﬂeet.
    Zurcher’s FLOW UTILITY OF RUNNING (i.e. not replacing) a bus is

    (1): $u_{1}(x_{1t},b)= θ_{0}+ θ_{1}x_{1t} + θ_{2}b$ [see assignment text]

    where x_{1t} is the mileage on the bus’s odometer (in 10,000s of miles) and b is a dummy variable indicating
    whether the bus is branded (meaning its manufacturer is high-end). The choice set is {0,1} where 0 denotes
    replacing the engine.
    Estimate the θ parameters assuming Zurcher is completely myopic.
    This amounts to estimating a simple binary logit model.
    (Note: you may estimate this any way you wish. I would recommend using the GLM package, but you may also
    use Optim with your own log likelihood function.)
    =#

    θ̂_glm = glm(@formula(Y ~ Odometer + Branded), df_long, Binomial(), LogitLink()) # GG: estimating using GLM package. Pooled binary logit using reshaped long data
    println(θ̂_glm)

#= GG: III - DYNAMIC ESTIMATION
Now I will walk you through how to estimate the DYNAMIC version of this model using BACKWARDS RECURSION.
With discount factor β, the DIFFERENTIATED CONDITIONAL VALUE FUNCTION for running the bus (relative to replacing it) is

(2): $ v_{1t}(x_{t},b) - v_{0t}(x_{t},b) =
        θ_{0} + θ_{1}x_{1t} + θ_{2}b                    [GG: same as static so far, flow utility]
        + β ∫ V_{t+1}(x_{t+1}, b) dF(x_{t+1} | x_{t})     [GG: recursive component]
        $

where V_{t+1} is the VALUE FUNCTION and the integral is over TRANSITIONS in the mileage states x_{t}.

We will APPROXIMATE the INTEGRAL with a SUMMATION, which means that we will specify a DISCRETE MASS FUNCTION for f(x_{t+1}  |x_{t}).
This probability mass function depends on the current odometer reading (x_{1t}),
whether the engine is newly replaced (i.e. d_{t−1}=0),
and on the value of another state variable x_2 which measures the usage intensity of the bus’s route
(i.e. high values of x_2 imply a low usage intensity and vice versa).
We discretize the mileage transitions into 1,250-mile bins (i.e. 0.125 units of x-{1t}).
We specify x_2 as a discrete uniform distribution ranging from 0.25 to 1.25 with 0.01 unit increments.

Formally, we are DISCRETELY (but not discreetly!) APPRXIMATING an EXPONENTIAL DISTRIBUTION:

(3): [see assignment text]

You will not need to program (3); I will provide code for this part. Under this formulation, (2) can be written as:

(4): [see assignment text]

Finally, we can simplify (4) since we know that V_{t+1 = log( ∑_{k} exp(v_{k,t+1}) ) when we assume
that unobserved utility is drawn from a T1EV distribution (as we do here):

(5): [see assignment text]

Estimation of our dynamic model now requires two steps:

#1 SOLVING THE MODEL
First, we need to solve the value functions for a given value of our parameters θ.
The way we do this is by BACKWARDS RECURSION. We know that V_{t+1} = 0 in our ﬁnal period (i.e. when t = T).
Then we work backwards to obtain the future value at every possible state in our model.
This will include many states that do not actually show up in our data.

# 2 ESTIMATING THE MODEL
Second, once we’ve solved the value functions, we use maximum likelihood to estimate the parameters θ.
The log likelihood function in this case is simply:
(6)-(7): [see assignment text]

Now estimate the θ’s assuming that Zurcher discounts the future with discount factor β = 0.9.
I will walk you through speciﬁc steps for how to do this
=#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3a: Read in the data for the dynamic model.
    # This can be found at the same URL as listed at the top of p. 2, but remove the "Beta0" from the CSV ﬁlename.
    # Rather than reshaping the data to “long” format as in question 1, we want to keep the data in “wide” format.
    # Thus, columns :Y1 through :Y20 should be converted to an array labeled Y which has
    # dimension 1000 × 20 where N = 1000 and T = 20. And similarly for columns starting with :Odo and :Xst.
    # Variables :Xst* and :Zst keep track of which discrete bin of the f j ’s the given observation falls into.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # [ Load in the data ]
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdata.csv"
    df = CSV.read(HTTP.get(url).body)
    Y = Matrix(df[:,[:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20]])
    X = Matrix(df[:,[:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20]])
    Z = Vector(df[:,:RouteUsage])
    B = Vector(df[:,:Branded])
    N = size(Y,1)
    T = size(Y,2)
    Xstate = Matrix(df[:,[:Xst1,:Xst2,:Xst3,:Xst4,:Xst5,:Xst6,:Xst7,:Xst8,:Xst9,:Xst10,:Xst11,:Xst12,:Xst13,:Xst14,:Xst15,:Xst16,:Xst17,:Xst18,:Xst19,:Xst20]])
    Zstate = Vector(df[:,:Zst])

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3b: Generate state transition matrices
    # Construct the state transition matrices, which are the f_j’s in (3). To do so, simply run the following code:
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    zval,zbin,xval,xbin,xtran = create_grids() # GG: creates grid points based on points given. See .jl file pf function
    # GG: returns bin numbers and associated values for z (rout usage) and x (odometer reading); and transition values for the X

    #=
    GG: zval and xval are the grids deﬁned at the bottom of p. 2, which respectively correspond to the route usage
    and odometer reading. zbin and xbin are the number of bins in zval and xval, respectively.
    xtran is a (zbin*xbin)×xbin Markov transition matrix that gives the probability of falling into each x_{1,t+1} bin
    given values of x_{1,t} and x_2 , according to the formula in (3).

    Note: A Markov transition matrix is a matrix where each row sums to 1 and moving from e.g. column 1 to column 4
    within a row gives the probability of moving from state 1 to state 4. Check out the Wikipedia page for more information
    =#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3c: Compute the future value terms for all possible states of the model
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # • First, initialize the future value array, which should be a 3-dimensional array of zeros.
    #   The size of the ﬁrst dimension should be the total number of grid points (i.e. the number of rows of xtran).
    #   The second dimension should be 2, which is the possible outcomes of :Branded.
    #   The third dimension should be T+1.
    # • Now write four nested for loops over each of the possible states:
    #   - Loop backwards over t from T+1 to 1
    #   – Loop over the two possible brand states {0,1}
    #   – Loop over the possible permanent route usage states (i.e. from 1 to zbin)
    #   – Loop over the possible odometer states (i.e. from 1 to xbin)
    # • Inside all of the for loops, make the following calculations
    #   – Create an object that marks the row of the transition matrix that we need to be looking at
    #     (based on the loop values of the two gridded state variables).
    #     This will be x + (z-1)*xbin (where x indexes the mileage bin and z indexes the route usage bin),
    #     given how the xtran matrix was constructed in the create grids() function.
    #   – Create the conditional value function for driving the bus (v_{1t}) based on the values of the state
    #     variables in the loop (not the values observed in the data). For example, for the mileage (x_{1t}),
    #     you should plug in xval[x] rather than :Odo. The difﬁcult part of the conditional value function is
    #     the discrete summation over the state transitions. For this, you need to grab the appropriate row
    #     (and all columns) of the xtran matrix, and then take the dot product with that and the all possible x_{1t}
    #     rows of the FV matrix. You should end up with something like xtran[row,:]'*FV[(z-1)*xbin+1:z*xbin,b+1,t+1]
    #     where b indexes the branded dummy and t indexes time periods.
    #   – Now create the conditional value function for replacing the engine (v_{0t}). For this, we repeat the same
    #     process as with v_{1t} except the θ’s are NORMALIZED to be 0. The code for the expected future value is
    #     the same as for v_{1t} with the exception that mileage resets to 0 after replacement, so instead of grabbing
    #     xtran[row,:] we want xtran[1+(z-1)*xbin,:].
    #   – Finally, update the future value array in period t by storing β log( exp(v_{0t}) + exp(v—1t}) ) in the t_th
    #     slice of the 3rd dimension of the array. This will be the new future value term for period t−1.
    #     Remember to set β = 0.9
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    @views @inbounds function likebus(α,β,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval)

        # FV in period T+1 = 0 since this is a finite horizon problem
        FV=zeros(zbin*xbin,2,T+1) # GG: initialize future values; initialize the future value array, which should be a 3-dimensional array of zeros.
        # GG: 2 columns is Branded=1 vs Branded=0

        # First loop: solve the backward recursion problem given the values of α
        # This will give the future value for *all* possible states that we might visit
        # This is why the FV array does not have an individual dimension
        for t=T:-1:1  # GG: Loop backwards over t from T+1 to 1; going backwards over time every period
            for b=0:1   # GG: Loop over the two possible brand states {0,1}; based on branded or not
                for z=1:zbin # GG: Loop over the possible permanent route usage states (i.e. from 1 to zbin); which bin am I in the x2 state variable
                    for x=1:xbin # GG: Loop over the possible odometer states (i.e. from 1 to xbin); which hin am I in the X1t state variable
                        # inputs to FV

                        # GG: object that marks the row of the transition matrix; to figurate the transition probabilities
                        row = x + (z-1)*xbin # which row of xtran should we look at? depends on milage bin x and route usage z

                        # GG: conditional value function for driving the bus (v_{1t})

                        v1  = α[1] + α[2]*xval[x] + α[3]*b + xtran[         row,:]⋅FV[(z-1)*xbin+1:z*xbin,b+1,t+1] # mileage bin is x, route usage z is permanent
                        # GG: flow utility component + summation over state transitions
                        # GG: set of probs xtran[row,:] DOT PRODUCT Future Values (FV)

                        # GG: conditional value function for replacing the engine (v_{0t})
                        v0  =                                xtran[1+(z-1)*xbin,:]⋅FV[(z-1)*xbin+1:z*xbin,b+1,t+1] # the engine got replaced => mileage is 0, so first bin

                        # FV is discounted log sum of the exponentiated conditional value functions #GG: Updating Future Value terms!!!
                        FV[row,b+1,t] = β*log(exp(v1) + exp(v0))
                        #GG: that's where the recursion comes in!
                    end
                end
            end
        end

        #:::::::::::::::::::::::::::::::::::::::::::::::::::
        # Question 3d: Construct the log likelihood using the future value terms from the previous step and only
        # using the observed states in the data. This will entail a for loop over buses and time periods.
        #:::::::::::::::::::::::::::::::::::::::::::::::::::
        # • Initialize the log likelihood value to be 0. (We will iteratively add to it as we loop over observations in the data)
        # • Create a variable that indexes the state transition matrix rows for the case where the bus has been replaced.
        #   This will be the same 1+(z-1)*xbin as in the conditional value function v 0t above. However, we need to plug in :Zst from the data rather than a hypothetical value z.
        # • Create a variable that indexes the state transition matrix rows for the case where the bus has not been replaced.
        #   This will be the same x + (z-1)*xbin as in v1t  above, except we substitute :Xst and :Zst for x and z.
        # • Now create the ﬂow utility component of v_{1t} − v_{0t} using the actual observed data on mileage and branding.
        # • Next, we need to add the appropriate discounted future value to round out our calculation of v_{1t} − v_{0t}.
        #   Here, we can difference the f_j’s as in (5). You should get something like (xtran[row1,:].-xtran[row0,:])'*FV[row0:row0+xbin-1,B[i]+1,t+1]
        # • Finally, create the choice probabilities for choosing each option as written in (7) and then create the
        #   log likelihood according to the summation in (6).

        # Second loop: form the likelihood given the future values implied by the previous α guess
        # Here, we will take the state-specifc FV's calculated in the first loop
        # But we will only use in the likelihood those state value that are actually visited
        like = 0 # GG: Initialize the log likelihood value to be 0.
        for i=1:N
            # GG: vector of indices in the ACTUAL DATA (no replacement)
            row0 = (Zstate[i]-1)*xbin+1 # this is the same argument as the index of xtran in v0 above, but we use the actual Z
            #z2   = (Zstate[i]-1)*xbin+1 # this is the same argument as the 1st part of the 1st index of the FV array above, but using actual Z
            #z3   = z2+xbin-1            # this is the same argument as the 2nd part of the 1st index of the FV array above, but using actual Z
            for t=1:T
                # GG: vector of indices in the ACTUAL DATA (replacement)
                row1  = Xstate[i,t] + (Zstate[i]-1)*xbin                                                                # this is the same as row in the first loop, except we use the actual X and Z
                v1    = α[1] + α[2]*X[i,t] + α[3]*B[i] + (xtran[row1,:].-xtran[row0,:])⋅FV[row0:row0+xbin-1,B[i]+1,t+1] # using the same formula as v1 above, except we use observed values of X and B, and we difference the transitions
                dem   = 1 + exp(v1)
                like -= ( (Y[i,t]==1)*v1 ) - log(dem) # negative of the binary logit likelihood
            end
        end
        return like
    end

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3e: Wrap all of the code you wrote in (c) and (d) into a function and set up the function so that it can be
    # passed to Optim. For example, you will need to return the negative of the log likelihood and you will need to have
    # the ﬁrst argument be the θ vector that we are trying to estimate
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: See row 213

    θ_start = rand(3)
    θ_true  = [2; -.15; 1]
    # how long to evaluate likelihood function once?
    println("Timing (twice) evaluation of the likelihood function")
    @time likebus(θ_start,0.9,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval)
    @time likebus(θ_start,0.9,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval)
    # estimate likelihood function
    θ̂_optim = optimize(a -> likebus(a,0.9,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval), θ_true, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true))
    θ̂_ddc = θ̂_optim.minimizer
    println(θ̂_ddc)

    return nothing
end

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3f: On the same line as the function, prepend the function declaration with the macros so that your code
    # says @views @inbounds function myfun() rather than function myfun(). This will give you more performant code.
    # On my machine, it cut the computation time in half.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: See row 213
    # @inbounds turns off dimensions checkin by Julia; consumes time, if your sure comformability is right you can turn it off to save time
    # @views treats items as immutable, makes things more efficiently
    # GG: You gotta make sure your functions work firt. Then you use these options to make things faster
    # GG: in Matlab it takes about 25mins; in Julia only 5mins; if we used the TwiceDifferentiable object, it'd be even faster

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3g: Wrap all of your code in an empty function as you’ve done with other problem sets
#:::::::::::::::::::::::::::::::::::::::::::::::::::
allwrap()

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3h: Try executing your script to estimate the likelihood function. This took about 4 minutes
# on my machine when I started from the estimates of the static model in Question 2.
#:::::::::::::::::::::::::::::::::::::::::::::::::::

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3i: Pat yourself on the back and grab a beverage of your choice, because that was a lot of work!
#:::::::::::::::::::::::::::::::::::::::::::::::::::


text:
using Distributions, RecursiveArrayTools, NLopt, Random
Random.seed!(123)

original_solution = VectorOfArray([(sol(t[i])) for i in 1:length(t)])
original_solution_matrix_form = convert(Array,original_solution)
weight = zeros(size(original_solution_matrix_form)[1],size(original_solution_matrix_form)[2])
error = zeros(size(original_solution_matrix_form)[1],size(original_solution_matrix_form)[2])

for i in 1:size(original_solution_matrix_form)[1]
 for j in 1:size(original_solution_matrix_form)[2]
  tmp = rand()
  weight[i,j] = tmp*tmp
  d = Normal(0,tmp)
  error[i,j] = rand(d)
 end
end

weighted_data = original_solution_matrix_form + error

weighted_cost_function = build_loss_objective(prob1,Tsit5(),
          L2Loss(t,weighted_data,data_weight=weight),maxiters=10000,verbose=false)
opt = Opt(:LN_COBYLA, 1)
min_objective!(opt, weighted_cost_function)
(minf,minx,ret) = NLopt.optimize(opt,[1.3])
@test minx[1] ≈ 1.5 atol=1e-2

ext: jl


text:
# Apply boundary conditions to the system

function apply_boundary_conditions_rhs_only(var, b, t)
    return apply_boundary_conditions_lhs_rhs(var, nothing, b, t)
end

function apply_boundary_conditions_lhs_rhs(var, A, b, t)
    if A === nothing
        rhs_only = true
    else
        rhs_only = false
    end
    multivar = typeof(var) <: Array;
    maxvarindex = 0;
    if multivar
        # multiple variables being solved for simultaneously
        dofs_per_node = 0;
        var_to_dofs = [];
        for vi=1:length(var)
            tmp = dofs_per_node;
            dofs_per_node += var[vi].total_components;
            push!(var_to_dofs, (tmp+1):dofs_per_node);
            maxvarindex = max(maxvarindex,var[vi].index);
        end
    else
        # one variable
        dofs_per_node = var.total_components;
        maxvarindex = var.index;
    end

    # If there is a reason to zero the rhs bdry entries, uncomment this.
    # But this destroys the case NO_BC.
    #b = zero_rhs_bdry_vals(b, dofs_per_node);

    bidcount = length(grid_data.bids); # the number of BIDs
    dirichlet_nodes = zeros(Int, 0); # Nodes corresponding to rows
    dirichlet_rows = zeros(Int, 0); # Dirichlet rows will be identity
    dirichlet_vals = zeros(0); # The values that go in RHS (b)
    neumann_nodes = zeros(Int, 0); # Nodes corresponding to rows
    neumann_rows = zeros(Int, 0); # Neumann rows will be set accordingly
    neumann_vals = zeros(0); # The values that go in RHS (b)
    if !rhs_only
        neumann_Is = zeros(Int,0);
        neumann_Js = zeros(Int,0);
        neumann_Vs = zeros(0);
    end
    if multivar
        dofind = 0;
        for vi=1:length(var)
            for compo=1:var[vi].total_components
                dofind = dofind + 1;
                for bid=1:bidcount
                    if prob.bc_type[var[vi].index, bid] == NO_BC
                        # do nothing
                    elseif prob.bc_type[var[vi].index, bid] == DIRICHLET
                        (tmprows, tmpnodes, tmpvals) = dirichlet_bc(prob.bc_func[var[vi].index, bid][compo], grid_data.bdryface[bid], t, dofind, dofs_per_node);
                        append!(dirichlet_rows, tmprows);
                        append!(dirichlet_vals, tmpvals);
                        append!(dirichlet_nodes, tmpnodes);
                    elseif prob.bc_type[var[vi].index, bid] == NEUMANN
                        if rhs_only
                            (tmprows, tmpnodes, tmpvals) = neumann_bc(prob.bc_func[var[vi].index, bid][compo], grid_data.bdryface[bid], bid, t, dofind, dofs_per_node, rhs_only=true);
                            append!(neumann_rows, tmprows);
                            append!(neumann_vals, tmpvals);
                            append!(neumann_nodes, tmpnodes);
                        else
                            (tmprows, tmpnodes, tmpvals, tmpIs, tmpJs, tmpVs) = neumann_bc(prob.bc_func[var[vi].index, bid][compo], grid_data.bdryface[bid], bid, t, dofind, dofs_per_node);
                            append!(neumann_rows, tmprows);
                            append!(neumann_vals, tmpvals);
                            append!(neumann_nodes, tmpnodes);
                            append!(neumann_Is, tmpIs);
                            append!(neumann_Js, tmpJs);
                            append!(neumann_Vs, tmpVs);
                        end

                    elseif prob.bc_type[var[vi].index, bid] == ROBIN
                        printerr("Robin BCs not ready.");
                    else
                        printerr("Unsupported boundary condition type: "*prob.bc_type[var[vi].index, bid]);
                    end
                end
            end
        end
    else
        for d=1:dofs_per_node
            dofind = d;
            for bid=1:bidcount
                if prob.bc_type[var.index, bid] == NO_BC
                    # do nothing
                elseif prob.bc_type[var.index, bid] == DIRICHLET
                    (tmprows, tmpnodes, tmpvals) = dirichlet_bc(prob.bc_func[var.index, bid][d], grid_data.bdryface[bid], t, dofind, dofs_per_node);
                    append!(dirichlet_rows, tmprows);
                    append!(dirichlet_vals, tmpvals);
                    append!(dirichlet_nodes, tmpnodes);
                elseif prob.bc_type[var.index, bid] == NEUMANN
                    if rhs_only
                        (tmprows, tmpnodes, tmpvals) = neumann_bc(prob.bc_func[var.index, bid][d], grid_data.bdryface[bid], bid, t, dofind, dofs_per_node, rhs_only=true);
                        append!(neumann_rows, tmprows);
                        append!(neumann_vals, tmpvals);
                        append!(neumann_nodes, tmpnodes);
                    else
                        (tmprows, tmpnodes, tmpvals, tmpIs, tmpJs, tmpVs) = neumann_bc(prob.bc_func[var.index, bid][d], grid_data.bdryface[bid], bid, t, dofind, dofs_per_node);
                        append!(neumann_rows, tmprows);
                        append!(neumann_vals, tmpvals);
                        append!(neumann_nodes, tmpnodes);
                        append!(neumann_Is, tmpIs);
                        append!(neumann_Js, tmpJs);
                        append!(neumann_Vs, tmpVs);
                    end
                elseif prob.bc_type[var.index, bid] == ROBIN
                    printerr("Robin BCs not ready.");
                else
                    printerr("Unsupported boundary condition type: "*prob.bc_type[var.index, bid]);
                end
            end
        end
    end

    ## Now I have all of the rows, nodes and vals
    # First apply Neumann, then Dirichlet(can overwrite Neumann), then zero parts that are for non-owned nodes if needed.
    if !rhs_only
        if length(neumann_Is)>0
            A = insert_sparse_rows(A, neumann_Is, neumann_Js, neumann_Vs);
        end
        if length(dirichlet_rows)>0
            A = identity_rows(A, dirichlet_rows, length(b));
        end
    end
    for i=1:length(neumann_rows)
        b[neumann_rows[i]] = neumann_vals[i];
    end
    for i=1:length(dirichlet_rows)
        b[dirichlet_rows[i]] = dirichlet_vals[i];
    end
    # For partitioned meshes only
    if config.num_partitions > 1
        rows_to_zero = zeros(Int,0);
        # Zero all entries for boundary nodes that are not owned
        # Base this on the global boundary flags, not local because of annoying cases
        for ni=1:length(grid_data.partition2global)
            if grid_data.global_bdry_index[grid_data.partition2global[ni]] > 0 && grid_data.node_owner[ni] != config.partition_index
                if dofs_per_node == 1
                    b[ni] = 0;
                    push!(rows_to_zero, ni);
                else
                    for di=1:dofs_per_node
                        rowid = (ni-1)*dofs_per_node+di;
                        b[rowid] = 0;
                        push!(rows_to_zero, rowid);
                    end
                end
            end
        end

        if !rhs_only && length(rows_to_zero) > 0
            A = identity_rows(A, rows_to_zero, length(b), diagonal_val=0.0);
        end
    end

    # Reference points
    if size(prob.ref_point,1) >= maxvarindex
        if multivar
            posind = zeros(Int,0);
            vals = zeros(0);
            for vi=1:length(var)
                if prob.ref_point[var[vi].index,1]
                    eii = prob.ref_point[var[vi].index, 2];
                    tmp = (grid_data.glbvertex[eii[1], eii[2]] - 1)*dofs_per_node + var_to_dofs[vi][1];
                    if length(prob.ref_point[var[vi].index, 3]) > 1
                        tmp = tmp:(tmp+length(prob.ref_point[var[vi].index, 3])-1);
                    end
                    posind = [posind; tmp];
                    vals = [vals; prob.ref_point[var[vi].index, 3]];
                end
            end
            if length(vals) > 0
                if !rhs_only A = identity_rows(A, posind, length(b)); end
                b[posind] = vals;
            end

        else
            if prob.ref_point[var.index,1]
                eii = prob.ref_point[var.index, 2];
                posind = (grid_data.glbvertex[eii[1], eii[2]] - 1)*dofs_per_node + 1;
                if length(prob.ref_point[var.index, 3]) > 1
                    posind = posind:(posind+length(prob.ref_point[var[vi].index, 3])-1);
                else
                    posind = [posind];
                end
                if !rhs_only A = identity_rows(A, posind, length(b)); end
                b[posind] = prob.ref_point[var.index, 3];
            end
        end
    end

    if rhs_only
        return b;
    else
        return (A, b);
    end
end

# This is just to speed things up (huge improvement)
function is_in_rows(k, rows)
    # rows are sorted
    b = 1;
    e = length(rows);
    c = Int(round((e+1)/2));
    while e-b > 1
        if k == rows[c]
            return true;
        elseif k > rows[c]
            b=c;
        else
            e=c;
        end
        c = Int(round((b+e)/2));
    end
    return k==rows[b] || k==rows[e];
end

# zeros the rows and puts a 1 on the diagonal
function identity_rows(A, urows, N; diagonal_val=1)
    if length(urows) < 1
        return A;
    end
    if issparse(A)
        (I, J, V) = findnz(A);
        eN = length(I);
        newI = zeros(Int,eN);
        newJ = zeros(Int,eN);
        newV = zeros(eN);
        newind = 0;
        (M,N) = size(A);

        rows = sort(urows);

        # remove duplicates
        newrows = zeros(Int, length(rows));
        newcount = 1;
        newrows[1] = rows[1];
        for i=2:length(rows)
            if rows[i] != rows[i-1] # because sorted
                newcount += 1;
                newrows[newcount] = rows[i];
            end
        end
        rows = newrows[1:newcount];

        # Remove bdry row elements
        for k in 1:length(I)
            isbdry = is_in_rows(I[k], rows);
            if !isbdry
                newind = newind+1;
                newI[newind] = I[k];
                newJ[newind] = J[k];
                newV[newind] = V[k];
            end
        end

        # Add diogonals
        rn = length(rows);
        if newind + rn <= eN
            for i=1:rn
                newind = newind+1;
                newI[newind] = rows[i];
                newJ[newind] = rows[i];
                newV[newind] = diagonal_val;
            end
            newI = newI[1:newind];
            newJ = newJ[1:newind];
            newV = newV[1:newind];
        else
            toadd = rn - (eN-newind);
            for i=1:(eN-newind)
                newind = newind+1;
                newI[newind] = rows[i];
                newJ[newind] = rows[i];
                newV[newind] = diagonal_val;
            end
            append!(newI, rows[toadd:rn]);
            append!(newJ, rows[toadd:rn]);
            append!(newV, diagonal_val * ones(toadd));
        end

        return sparse(newI,newJ,newV, M,N);

    else
        for i=1:length(rows)
            A[rows[i],:] .= 0;
            A[rows[i],rows[i]] = diagonal_val;
        end
        return A;
    end
end

# Inserts the rows of S in A
function insert_rows(A, S, rows, N)
    if issparse(A)
        (I, J, V) = findnz(A);
        (M,N) = size(A);
        # determine which elements to remove
        toskip = zeros(Bool, length(I));
        includen = 0;
        for k = 1:length(I)
            for i=1:length(rows)
                if I[k] == rows[i]
                    toskip[k] = 1;
                end
            end
            if toskip[k] == 0
                includen = includen + 1;
            end
        end

        # put the remaining elements in a new matrix
        newI = zeros(Int, includen);
        newJ = zeros(Int, includen);
        newV = zeros(includen);
        ind = 1;
        for k = 1:length(I)
            if toskip[k] == 0
                newI[ind] = I[k];
                newJ[ind] = J[k];
                newV[ind] = V[k];
                ind = ind + 1;
            end
        end

        # Do something similar to extract elements of S
        (I, J, V) = findnz(S);

        toskip = ones(Int, length(I));
        includen = 0;
        for k = 1:length(I)
            for i=1:length(rows)
                if I[k] == rows[i]
                    toskip[k] = 0;
                end
            end
            if toskip[k] == 0
                includen = includen + 1;
            end
        end
        newI2 = zeros(Int, includen);
        newJ2 = zeros(Int, includen);
        newV2 = zeros(includen);
        ind = 1;
        for k = 1:length(I)
            if toskip[k] == 0
                newI2[ind] = I[k];
                newJ2[ind] = J[k];
                newV2[ind] = V[k];
                ind = ind + 1;
            end
        end

        append!(newI, newI2);
        append!(newJ, newJ2);
        append!(newV, newV2);
        return sparse(newI,newJ,newV, M, N);
    else
        for i=1:length(rows)
            A[rows[i],:] = S[rows[i],:];
        end
        return A;
    end
end

function insert_sparse_rows(A, SI, SJ, SV)
    (I, J, V) = findnz(A);
    (M,N) = size(A);

    # since we need to sort the three arrays in sync, just get the sorted permutation
    sorted_order = sortperm(SI);

    SI = SI[sorted_order];
    SJ = SJ[sorted_order];
    SV = SV[sorted_order];

    # Zero existing elements in SI rows
    for k in 1:length(I)
        if is_in_rows(I[k], SI)
            V[k] = 0;
        end
    end

    # append S values
    append!(I, SI);
    append!(J, SJ);
    append!(V, SV);

    # Form sparse matrix
    return sparse(I, J, V, M, N);
end

function zero_rhs_bdry_vals(rhs, dofs)
    for bid=1:length(grid_data.bids)
        for fi=1:length(grid_data.bdryface[bid])
            face_nodes = grid_data.face2glb[:, 1, grid_data.bdryface[bid][fi]];
            for ni=1:length(face_nodes)
                rhs_ind = ((face_nodes[ni]-1)*dofs + 1):(face_nodes[ni]*dofs);
                rhs[rhs_ind] .= 0;
            end
        end
    end
    return rhs;
end

function copy_bdry_vals_to_variables(var, vec, grid, dofs_per_node; zero_vals=true)
    if typeof(var) <: Array
        dofind = 0;
        for vi=1:length(var)
            for compo=1:length(var[vi].symvar)
                dofind = dofind + 1;
                for bid=1:size(prob.bc_type,2)
                    if prob.bc_type[var[vi].index, bid] == DIRICHLET
                        for i = 1:length(grid.bdryface[bid]) # loop over faces with this BID
                            fid = grid.bdryface[bid][i];
                            face_nodes = grid.face2glb[:,1,fid];
                            for ni=1:length(face_nodes)
                                node = face_nodes[ni];
                                var[vi].values[compo, node] = vec[(node-1)*dofs_per_node + dofind];
                                if zero_vals
                                    vec[(node-1)*dofs_per_node + dofind] = 0;
                                end
                            end
                        end
                    end
                end
            end
        end
    else
        for d=1:dofs_per_node
            dofind = d;
            for bid=1:size(prob.bc_type,2)
                if prob.bc_type[var.index, bid] == DIRICHLET
                    for i = 1:length(grid.bdryface[bid]) # loop over faces with this BID
                        fid = grid.bdryface[bid][i];
                        face_nodes = grid.face2glb[:,1,fid];
                        for ni=1:length(face_nodes)
                            node = face_nodes[ni];
                            var.values[d, node] = vec[(node-1)*dofs_per_node + dofind];
                            if zero_vals
                                vec[(node-1)*dofs_per_node + dofind] = 0;
                            end
                        end
                    end
                end
            end
        end
    end
end

function copy_bdry_vals_to_vector(var, vec, grid, dofs_per_node)
    if typeof(var) <: Array
        dofind = 0;
        for vi=1:length(var)
            for compo=1:length(var[vi].symvar)
                dofind = dofind + 1;
                for bid=1:size(prob.bc_type,2)
                    if prob.bc_type[var[vi].index, bid] == DIRICHLET
                        for i = 1:length(grid.bdryface[bid]) # loop over faces with this BID
                            fid = grid.bdryface[bid][i];
                            face_nodes = grid.face2glb[:,1,fid];
                            for ni=1:length(face_nodes)
                                node = face_nodes[ni];
                                vec[(node-1)*dofs_per_node + dofind] = var[vi].values[compo, node];
                            end
                        end
                    end
                end
            end
        end
    else
        for d=1:dofs_per_node
            dofind = d;
            for bid=1:size(prob.bc_type,2)
                if prob.bc_type[var.index, bid] == DIRICHLET
                    for i = 1:length(grid.bdryface[bid]) # loop over faces with this BID
                        fid = grid.bdryface[bid][i];
                        face_nodes = grid.face2glb[:,1,fid];
                        for ni=1:length(face_nodes)
                            node = face_nodes[ni];
                            vec[(node-1)*dofs_per_node + dofind] = var.values[d, node];
                        end
                    end
                end
            end
        end
    end
end

# This evaluates the BC at a specific node.
# That could mean:
# - the value of constant BCs
# - evaluate a genfunction.func for BCs defined by strings->genfunctions
# - evaluate a function for BCs defined by callback functions
function evaluate_at_nodes(val, nodes, face, t)
    N = length(nodes);
    dim = config.dimension;
    if typeof(val) <: Number
        result = fill(val, N);

    elseif typeof(val) == Coefficient && typeof(val.value[1]) == GenFunction
        result = zeros(N);
        if dim == 1
            for i=1:N
                result[i]=val.value[1].func(grid_data.allnodes[1,nodes[i]],0,0,t,nodes[i],face);
            end
        elseif dim == 2
            for i=1:N
                result[i]=val.value[1].func(grid_data.allnodes[1,nodes[i]],grid_data.allnodes[2,nodes[i]],0,t,nodes[i],face);
            end
        else
            for i=1:N
                result[i]=val.value[1].func(grid_data.allnodes[1,nodes[i]],grid_data.allnodes[2,nodes[i]],grid_data.allnodes[3,nodes[i]],t,nodes[i],face);
            end
        end

    elseif typeof(val) == GenFunction
        result = zeros(N);
        if dim == 1
            for i=1:N
                result[i]=val.func(grid_data.allnodes[1,nodes[i]],0,0,t,nodes[i],face);
            end
        elseif dim == 2
            for i=1:N
                result[i]=val.func(grid_data.allnodes[1,nodes[i]],grid_data.allnodes[2,nodes[i]],0,t,nodes[i],face);
            end
        else
            for i=1:N
                result[i]=val.func(grid_data.allnodes[1,nodes[i]],grid_data.allnodes[2,nodes[i]],grid_data.allnodes[3,nodes[i]],t,nodes[i],face);
            end
        end

    elseif typeof(val) == CallbackFunction
        result = zeros(N);
        for i=1:N
            #form a dict for the arguments. x,y,z,t are always included
            arg_list = [];
            append!(arg_list, [("x", grid_data.allnodes[1,nodes[i]]),
                        ("y", dim>1 ? grid_data.allnodes[2,nodes[i]] : 0),
                        ("z", dim>2 ? grid_data.allnodes[3,nodes[i]] : 0),
                        ("t", t)]);
            # Add in the requires list
            for r in val.args
                foundit = false;
                # is it an entity?
                if typeof(r) == Variable
                    foundit = true;
                    if size(r.values,1) == 1
                        push!(arg_list, (string(r.symbol), r.values[1,nodes[i]]));
                    else
                        push!(arg_list, (string(r.symbol), r.values[:,nodes[i]]));
                    end
                elseif typeof(r) == Coefficient
                    # TODO: evaluate coefficient at node
                    foundit = true;
                    push!(arg_list, (string(r.symbol), evaluate_at_nodes(r, nodes[i], face, t)));
                elseif typeof(r) == String
                    # This could also be a variable, coefficient, or special thing like normal,
                    if r in ["x","y","z","t"]
                        foundit = true;
                        # These are already included
                    elseif r == "normal"
                        foundit = true;
                        push!(arg_list, ("normal", grid_data.facenormals[:,face]));
                    else
                        for v in variables
                            if string(v.symbol) == r
                                foundit = true;
                                if size(v.values,1) == 1
                                    push!(arg_list, (r, v.values[1,nodes[i]]));
                                else
                                    push!(arg_list, (r, v.values[:,nodes[i]]));
                                end
                                break;
                            end
                        end
                        for c in coefficients
                            if string(c.symbol) == r
                                foundit = true;
                                push!(arg_list, (r, evaluate_at_nodes(c, nodes[i], face, t)));
                                break;
                            end
                        end
                    end
                else
                    # What else could it be?
                end

                if !foundit
                    # Didn't figure this thing out.
                    printerr("Unknown requirement for callback function "*string(fun)*" : "*string(requires[i]));
                end
            end

            # Build the dict
            args = Dict(arg_list);

            # call the function
            result[i] = val.func(args);
        end

    end
    return result;
end

function dirichlet_bc(val, bdryface, t=0, dofind = 1, totaldofs = 1)
    # Assuming every face has the same number of nodes
    nodes_per_face = size(grid_data.face2glb,1);
    nnodes = length(bdryface)*nodes_per_face;
    bdry_rows = zeros(Int, nnodes);
    bdry_nodes = zeros(Int, nnodes);
    bdry_vals = zeros(nnodes);

    for fi = 1:length(bdryface)
        if bdryface[fi] < 1
            println("found bad bdryface: "*string(bdryface));
        end
        # global indices for the nodes on the face
        face_nodes = grid_data.face2glb[:, 1, bdryface[fi]];

        # offset for multi dof
        if totaldofs > 1
            offsetglb = (face_nodes.-1) .* totaldofs .+ dofind;
        else
            offsetglb = face_nodes;
        end

        bdry_rows[((fi-1)*nodes_per_face+1):(fi*nodes_per_face)] = offsetglb;
        bdry_nodes[((fi-1)*nodes_per_face+1):(fi*nodes_per_face)] = face_nodes;
        bdry_vals[((fi-1)*nodes_per_face+1):(fi*nodes_per_face)] = evaluate_at_nodes(val, face_nodes, bdryface[fi], t);
    end

    return (bdry_rows, bdry_nodes, bdry_vals);
end

function neumann_bc(val, bdryface, bid, t=0, dofind = 1, totaldofs = 1; rhs_only=false)
    # Assuming every element has the same number of nodes
    nodes_per_element = size(grid_data.loc2glb,1);
    nodes_per_face = size(grid_data.face2glb,1);
    nnodes = length(bdryface)*nodes_per_face;
    bdry_rows = zeros(Int, nnodes);
    bdry_nodes = zeros(Int, nnodes);
    bdry_vals = zeros(nnodes);

    # To deal with the possibility of overlapping nodes between elements or other bids
    # limit it to nodes assigned to this BID and assign each node to one of the bdryfaces.
    # Need to form a list of valid nodes for each face in bdryface.
    bidnodes = grid_data.bdry[bid];
    node_available = fill(true, length(bidnodes));
    valid_face_mask = zeros(Bool, nodes_per_face, length(bdryface));
    for fi=1:length(bdryface)
        face_nodes = grid_data.face2glb[:, 1, bdryface[fi]];
        for ni=1:length(face_nodes)
            bnode_ind = indexin([face_nodes[ni]], bidnodes);
            if !(bnode_ind[1] === nothing)
                if node_available[bnode_ind[1]]
                    # This is a valid node for this face. Reserve it
                    node_available[bnode_ind[1]] = false;
                    valid_face_mask[ni,fi] = true;
                end
            end
        end
    end
    total_entries = length(bidnodes) * nodes_per_element;

    if !rhs_only
        # These will be sparse rows to insert into A
        S_I = zeros(Int, total_entries);
        S_J = zeros(Int, total_entries);
        S_V = zeros(total_entries);
        next_S_ind = 1;

        # This can be precomputed for uniform grid meshes
        if config.mesh_type == UNIFORM_GRID && config.geometry == SQUARE
            precomputed = true;
            glb = grid_data.loc2glb[:,1];
            detJ = geo_factors.detJ[1];
            J = geo_factors.J[1]
            if config.dimension == 1
                RD1 = zeros(size(refel.Ddr));
                # Multiply rows of Qr and Ddr by J.rx
                for i=1:size(RD1,2) # loop over columns
                    RD1[:,i] = J.rx .* refel.Ddr[:,i];
                end

            elseif config.dimension == 2
                RD1 = zeros(size(refel.Ddr));
                RD2 = zeros(size(refel.Ddr));
                for i=1:size(RD1,2)
                    RD1[:,i] = J.rx .* refel.Ddr[:,i] + J.sx .* refel.Dds[:,i];
                    RD2[:,i] = J.ry .* refel.Ddr[:,i] + J.sy .* refel.Dds[:,i];
                end

            elseif config.dimension == 3
                RD1 = zeros(size(refel.Ddr));
                RD2 = zeros(size(refel.Ddr));
                RD3 = zeros(size(refel.Ddr));
                for i=1:size(RD1,2)
                    RD1[:,i] = J.rx .* refel.Ddr[:,i] + J.sx .* refel.Dds[:,i] + J.tx .* refel.Ddt[:,i];
                    RD2[:,i] = J.ry .* refel.Ddr[:,i] + J.sy .* refel.Dds[:,i] + J.ty .* refel.Ddt[:,i];
                    RD3[:,i] = J.rz .* refel.Ddr[:,i] + J.sz .* refel.Dds[:,i] + J.tz .* refel.Ddt[:,i];
                end
            end
        else
            precomputed = false;
        end
    end

    next_row_index = 1;
    for fi = 1:length(bdryface)
        # Local indices for the nodes on the face
        flocal = refel.face2local[grid_data.faceRefelInd[1,bdryface[fi]]];

        # global indices for the nodes on the face
        face_nodes = grid_data.face2glb[:, 1, bdryface[fi]];

        # remove nodes that do not belong to this face and bid
        newflocal = zeros(Int, length(flocal));
        newface_nodes = zeros(Int, length(face_nodes));
        node_count = 0;
        for ni=1:length(face_nodes)
            if valid_face_mask[ni, fi]
                node_count += 1;
                newflocal[node_count] = flocal[ni];
                newface_nodes[node_count] = face_nodes[ni];
            end
        end
        face_nodes = newface_nodes[1:node_count];
        flocal = newflocal[1:node_count];

        # offset for multi dof
        if totaldofs > 1
            offsetface = (face_nodes.-1) .* totaldofs .+ dofind;
        else
            offsetface = face_nodes;
        end

        if !rhs_only
            # find the relevant element
            e = grid_data.face2element[1,bdryface[fi]] # >0 ? grid_data.face2element[1,bdryface[fi]] : grid_data.face2element[2,bdryface[fi]];
            glb = grid_data.loc2glb[:,e];       # global indices of this element's nodes
            # offset for multi dof
            if totaldofs > 1
                elglb = (glb.-1) .* totaldofs .+ dofind;
            else
                elglb = glb;
            end

            # Build diff matrices if needed
            if !precomputed
                detJ = geo_factors.detJ[e];
                J = geo_factors.J[e]
                if config.dimension == 1
                    RD1 = zeros(size(refel.Ddr));
                    # Multiply rows of Qr and Ddr by J.rx
                    for i=1:size(RD1,2) # loop over columns
                        RD1[:,i] = J.rx .* refel.Ddr[:,i];
                    end

                elseif config.dimension == 2
                    RD1 = zeros(size(refel.Ddr));
                    RD2 = zeros(size(refel.Ddr));
                    for i=1:size(RD1,2)
                        RD1[:,i] = J.rx .* refel.Ddr[:,i] + J.sx .* refel.Dds[:,i];
                        RD2[:,i] = J.ry .* refel.Ddr[:,i] + J.sy .* refel.Dds[:,i];
                    end

                elseif config.dimension == 3
                    RD1 = zeros(size(refel.Ddr));
                    RD2 = zeros(size(refel.Ddr));
                    RD3 = zeros(size(refel.Ddr));
                    for i=1:size(RD1,2)
                        RD1[:,i] = J.rx .* refel.Ddr[:,i] + J.sx .* refel.Dds[:,i] + J.tx .* refel.Ddt[:,i];
                        RD2[:,i] = J.ry .* refel.Ddr[:,i] + J.sy .* refel.Dds[:,i] + J.ty .* refel.Ddt[:,i];
                        RD3[:,i] = J.rz .* refel.Ddr[:,i] + J.sz .* refel.Dds[:,i] + J.tz .* refel.Ddt[:,i];
                    end
                end
            end

            # Add to S matrix
            for i=1:length(elglb)
                irange = (next_S_ind + (i-1)*node_count):(next_S_ind + i*node_count - 1);
                S_I[irange] = elglb[flocal];
                S_J[irange] .= elglb[i];
            end
            if config.dimension == 1
                S1_V= RD1[flocal,:][:];

            elseif config.dimension == 2
                S1_V= RD1[flocal,:][:];
                S2_V= RD2[flocal,:][:];

            elseif config.dimension == 3
                S1_V= RD1[flocal,:][:];
                S2_V= RD2[flocal,:][:];
                S3_V= RD3[flocal,:][:];
            end

            # Add the right components of S1,S2,S3 according to normal vector
            ind_range = next_S_ind:(next_S_ind + node_count * nodes_per_element - 1);
            norm = grid_data.facenormals[:, bdryface[fi]];
            # println("bid "*string(bid)*" normal "*string(norm));
            if config.dimension == 1
                S_V[ind_range] = norm[1] .* S1_V;
            elseif config.dimension == 2
                S_V[ind_range] = norm[1] .* S1_V + norm[2] .* S2_V;
            elseif config.dimension == 3
                S_V[ind_range] = norm[1] .* S1_V + norm[2] .* S2_V + norm[3] .* S3_V;
            end

            # Increase S index
            next_S_ind += node_count * nodes_per_element;
        end

        bdry_nodes[next_row_index:(next_row_index+length(face_nodes)-1)] = face_nodes;
        bdry_rows[next_row_index:(next_row_index+length(face_nodes)-1)] = offsetface;
        bdry_vals[next_row_index:(next_row_index+length(face_nodes)-1)] = evaluate_at_nodes(val, face_nodes, bdryface[fi], t);
        next_row_index += length(face_nodes);
    end

    bdry_nodes = bdry_nodes[1:(next_row_index-1)];
    bdry_rows = bdry_nodes[1:(next_row_index-1)];
    bdry_vals = bdry_nodes[1:(next_row_index-1)];

    if rhs_only
        return (bdry_rows, bdry_nodes, bdry_vals);
    else
        if next_S_ind <= length(S_I)
            # This should not happen
            log_entry("Unexpected entry count when doing Neumann BC. expected "*string(length(S_I))*" found "*string(next_S_ind-1));
            S_I = S_I[1:next_S_ind-1];
            S_J = S_J[1:next_S_ind-1];
            S_V = S_V[1:next_S_ind-1];
        end
        return (bdry_rows, bdry_nodes, bdry_vals, S_I, S_J, S_V);
    end

end

max_stars_repo_name: paralab/Finch

max_issues_repo_name: paralab/Finch

max_forks_repo_name: paralab/Finch


text:
struct QuickMDP{ID,S,A,D<:NamedTuple} <: MDP{S,A}
    data::D
end

"""
    QuickMDP(gen::Function, [id]; kwargs...)

Construct a generative MDP model with the function `gen` and keyword arguments.

`gen` should take three arguments: a state, an action, and a random number generator. It should return a `NamedTuple` with keys `sp` for the next state and `r` for the reward.

Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
QuickMDP(gen::Function, id=uuid4(); kwargs...) = QuickMDP(id; gen=gen, kwargs...)

"""
    QuickMDP([id]; kwargs...)

Construct an MDP model with keyword arguments. Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
function QuickMDP(id=uuid4(); kwargs...)
    kwd = Dict{Symbol, Any}(kwargs)

    for (k, v) in pairs(kwd)
        kwd[k] = preprocess(v)
    end
    quick_defaults!(kwd)

    S = infer_statetype(kwd)
    A = infer_actiontype(kwd)
    d = namedtuple(keys(kwd)...)(values(kwd)...)
    qm = QuickMDP{id, S, A, typeof(d)}(d)
    return qm
end

id(::QuickMDP{ID}) where ID = ID

struct QuickPOMDP{ID,S,A,O,D<:NamedTuple} <: POMDP{S,A,O}
    data::D
end

"""
    QuickPOMDP(gen::Function, [id]; kwargs...)

Construct a generative POMDP model with the function `gen` and keyword arguments.

`gen` should take three arguments: a state, an action, and a random number generator. It should return a `NamedTuple` with keys `sp` for the next state, `o` for the observation, and `r` for the reward.

Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
QuickPOMDP(gen::Function, id=uuid4(); kwargs...) = QuickPOMDP(id; gen=gen, kwargs...)

"""
    QuickPOMDP([id]; kwargs...)

Construct an POMDP model with keyword arguments. Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
function QuickPOMDP(id=uuid4(); kwargs...)
    kwd = Dict{Symbol, Any}(kwargs)

    for (k, v) in pairs(kwd)
        kwd[k] = preprocess(v)
    end
    quick_defaults!(kwd)
    quick_warnings(kwd)

    S = infer_statetype(kwd)
    A = infer_actiontype(kwd)
    O = infer_obstype(kwd)
    d = namedtuple(keys(kwd)...)(values(kwd)...)
    qm = QuickPOMDP{id, S, A, O, typeof(d)}(d)
    return qm
end

id(::QuickPOMDP{ID}) where ID = ID

const QuickModel = Union{QuickMDP, QuickPOMDP}

"Function that is called on each keyword argument before anything else is done. This was designed as a hook to allow other packages to handle PyObjects."
preprocess(x) = x

function quick_defaults!(kwd::Dict)
    kwd[:discount] = get(kwd, :discount, 1.0)
    kwd[:isterminal] = get(kwd, :isterminal, false)

    if !haskey(kwd, :stateindex)
        if haskey(kwd, :states)
            states = _call(Val(:states), kwd[:states], ())
            if hasmethod(length, typeof((states,))) && length(states) < Inf
                kwd[:stateindex] = Dict(s=>i for (i,s) in enumerate(states))
            end
        end
    end

    if !haskey(kwd, :actionindex)
        if haskey(kwd, :actions)
            actions = _call(Val(:actions), kwd[:actions], ())
            if hasmethod(length, typeof((actions,))) && length(actions) < Inf
                kwd[:actionindex] = Dict(s=>i for (i,s) in enumerate(actions))
            end
        end
    end

    if !haskey(kwd, :obsindex)
        if haskey(kwd, :observations)
            observations = _call(Val(:observations), kwd[:observations], ())
            if hasmethod(length, typeof((observations,))) && length(observations) < Inf
                kwd[:obsindex] = Dict(s=>i for (i,s) in enumerate(observations))
            end
        end
    end
end

function quick_warnings(kwd)
    if haskey(kwd, :initialstate)
        isd = _call(Val(:initialstate), kwd[:initialstate], ())
        try rand(MersenneTwister(0), isd)
        catch ex
            if ex isa MethodError || ex isa ArgumentError
                @warn("Unable to call rand(rng, $isd). Is the `initialstate` that you supplied a distribution?")
            else
                rethrow(ex)
            end
        end
    end

    if haskey(kwd, :reward) && !(kwd[:reward] isa Function)
        @warn("`reward` must be a function; got $(kwd[:reward])")
    end
end

function infer_statetype(kwd)
    if haskey(kwd, :statetype)
        st = _call(Val(:statetype), kwd[:statetype], (), NamedTuple())
    elseif haskey(kwd, :states)
        st = eltype(_call(Val(:states), kwd[:states], (), NamedTuple()))
    elseif haskey(kwd, :initialstate)
        st = typeof(rand(MersenneTwister(0), _call(Val(:initialstate), kwd[:initialstate], (), NamedTuple())))
    else
        st = Any
    end
    if st == Any
        @warn("Unable to infer state type for a Quick(PO)MDP; using Any. This may have significant performance consequences. Use the statetype keyword argument to specify a concrete state type.")
    end
    return st
end

function infer_actiontype(kwd)
    if haskey(kwd, :actiontype)
        at = _call(Val(:actiontype), kwd[:actiontype], (), NamedTuple())
    elseif haskey(kwd, :actions)
        at = eltype(_call(Val(:actions), kwd[:actions], (), NamedTuple()))
    else
        at = Any
    end
    if at == Any
        @warn("Unable to infer action type for a Quick(PO)MDP; using Any. This may have significant performance consequences. Use the actiontype keyword argument to specify a concrete action type.")
    end
    return at
end

function infer_obstype(kwd)
    if haskey(kwd, :obstype)
        ot = _call(Val(:obstype), kwd[:obstype], (), NamedTuple())
    elseif haskey(kwd, :observations)
        ot = eltype(_call(Val(:observations), kwd[:observations], (), NamedTuple()))
    elseif haskey(kwd, :initialobs) && haskey(kwd, :initialstate)
        s0 = rand(MersenneTwister(0), _call(Val(:initialstate), kwd[:initialstate], (), NamedTuple()))
        ot = typeof(rand(MersenneTwister(0), _call(Val(:initialobs), kwd[:initialobs], (s0,), NamedTuple())))
    else
        ot = Any
    end
    if ot == Any
        @warn("Unable to infer observation type for a QuickPOMDP; using Any. This may have significant performance consequences. Use the obstype keyword argument to specify a concrete observation type.")
    end
    return ot
end

function _call(namev::Val{name}, m::QuickModel, args, kwargs=NamedTuple()) where name
    _call(namev,
          get(m.data, name) do
              throw(MissingQuickArgument(m, name))
          end,
          args,
          kwargs)
end

_call(::Val, f::Function, args, kwargs=NamedTuple()) = f(args...; kwargs...)
_call(v::Val, object, args, kwargs=NamedTuple()) = object
_call(v::Val, d::Dict, args, kwargs=NamedTuple()) = d[args...]

macro forward_to_data(f)
    @assert f.head == :. "@forward_to_data must be used with a module-qualified function expression, e.g. @forward_to_data POMDPs.discount"
    quote
        $f(m::QuickModel, args...; kwargs...) = _call(Val($(f.args[2])), m, args, kwargs)
    end
end

function POMDPs.transition(m::QuickModel, s, a)
    if haskey(m.data, :transition)
        return m.data.transition(s, a)
    else
        throw(MissingQuickArgument(m, :transition, types=[Function], also=[:gen]))
    end
end

function POMDPs.observation(m::QuickPOMDP, args...)
    if haskey(m.data, :observation)
        obs = m.data[:observation]
        if static_hasmethod(obs, typeof(args))
            return obs(args...)
        elseif length(args) == 3 && static_hasmethod(obs, typeof(args[2:3]))
            return obs(args[2:3]...)
        else
            return obs(args...)
        end
        return m.data.observation(args...)
    else
        throw(MissingQuickArgument(m, :observation, types=[Function], also=[:gen]))
    end
end

function POMDPs.reward(m::QuickModel, args...)
    if haskey(m.data, :reward)
        r = m.data[:reward]
        if static_hasmethod(r, typeof(args)) # static_hasmethod could cause issues, but I think it is worth doing in this single spot
            return r(args...)
        elseif m isa POMDP && length(args) == 4
            if static_hasmethod(r, typeof(args[1:3])) # (s, a, sp, o) -> (s, a, sp)
                return r(args[1:3]...)
            elseif static_hasmethod(r, typeof(args[1:2])) # (s, a, sp, o) -> (s, a)
                return r(args[1:2]...)
            end
        elseif length(args) == 3 && static_hasmethod(r, typeof(args[1:2])) # (s, a, sp) -> (s, a)
            return r(args[1:2]...)
        else
            return r(args...)
        end
    else
        throw(MissingQuickArgument(m, :reward))
    end
end

@forward_to_data POMDPs.initialstate
@forward_to_data POMDPs.initialobs

function POMDPs.gen(m::QuickModel, s, a, rng)
    if haskey(m.data, :gen)
        return m.data.gen(s, a, rng)
    else
        return NamedTuple()
    end
end

@forward_to_data POMDPs.states
@forward_to_data POMDPs.actions
@forward_to_data POMDPs.observations
@forward_to_data POMDPs.discount

@forward_to_data POMDPs.stateindex
@forward_to_data POMDPs.actionindex
@forward_to_data POMDPs.obsindex
@forward_to_data POMDPs.isterminal

function POMDPModelTools.obs_weight(m::QuickPOMDP, args...)
    if haskey(m.data, :obs_weight)
        return _call(Val(:obs_weight), m, args)
    elseif haskey(m.data, :observation)
        return pdf(_call(Val(:observation), m, args[1:end-1]), args[end])
    else
        throw(MissingQuickArgument(m, :obs_weight, types[Function], also=[:observation]))
    end
end
@forward_to_data POMDPModelTools.render

function POMDPModelTools.StateActionReward(m::Union{QuickPOMDP,QuickMDP})
    if hasmethod(m.data[:reward], Tuple{statetype(m), actiontype(m)})
        return FunctionSAR(m)
    else
        return LazyCachedSAR(m)
    end
end


text:
module NBody

# include("librationpoints.jl")

end


text:
using Poirot, Test

@testset "Poirot" begin

@testset "Abstract" begin
  include("abstract.jl")
end

@testset "Compiler" begin
  include("compiler.jl")
end

@testset "Inference" begin
  include("inference.jl")
end

@testset "LAX" begin
  include("lax.jl")
end

end


text:
"""
    ReplicaState(v, w, pnorm, params, id)

Struct that holds all information needed for an independent run of the algorithm.

Can be advanced a step forward with [`advance!`](@ref).

# Fields

* `hamiltonian`: the model Hamiltonian.
* `v`: vector.
* `w`: working memory.
* `pnorm`: previous walker number (see [`walkernumber`](@ref)).
* `params`: the [`FCIQMCRunStrategy`](@ref).
* `id`: appended to reported columns.
"""
mutable struct ReplicaState{H,T,V,W,R<:FciqmcRunStrategy{T}}
    hamiltonian::H
    v::V       # vector
    w::W       # working memory. Maybe working memories could be shared among replicas?
    pnorm::T   # previous walker number - used to control the shift
    params::R  # params: step, laststep, dτ...
    id::String # id is appended to column names
end

function ReplicaState(h, v, w, params, id="")
    if isnothing(w)
        w = similar(v)
    end
    return ReplicaState(h, v, w, walkernumber(v), params, id)
end

function Base.show(io::IO, r::ReplicaState)
    print(
        io,
        "ReplicaState(v: ", length(r.v), "-element ", nameof(typeof(r.v)),
        ", w: ", length(r.w), "-element ", nameof(typeof(r.w)), ")"
    )
end

"""
    QMCState

Holds all information needed to run FCIQMC, except the data frame. Holds a `NTuple` of
`ReplicaState`s and various strategies that control the algorithm.
"""
struct QMCState{
    H,
    N,
    R<:ReplicaState,
    MS<:MemoryStrategy,
    RS<:ReportingStrategy,
    SS<:ShiftStrategy,
    TS<:TimeStepStrategy,
    TH<:ThreadingStrategy,
    RRS<:ReplicaStrategy,
    PS<:NTuple{<:Any,PostStepStrategy},
}
    hamiltonian::H
    replicas::NTuple{N,R}
    maxlength::Ref{Int}

    m_strat::MS
    r_strat::RS
    s_strat::SS
    τ_strat::TS
    threading::TH
    post_step::PS
    replica::RRS
end

"""
    _n_walkers(v, s_strat)
Returns an estimate of the expected number of walkers as an integer.
"""
function _n_walkers(v, s_strat)
    n = if hasfield(typeof(s_strat), :targetwalkers)
        s_strat.targetwalkers
    else # e.g. for LogUpdate()
        walkernumber(v)
    end
    return ceil(Int, max(real(n), imag(n)))
end

function QMCState(
    hamiltonian, v;
    laststep=nothing,
    dτ=nothing,
    wm=nothing,
    params::FciqmcRunStrategy=RunTillLastStep(
        laststep = 100,
        shift = float(valtype(v))(diagonal_element(
            hamiltonian,
            starting_address(hamiltonian)
        ))
    ),
    s_strat::ShiftStrategy=DoubleLogUpdate(),
    r_strat::ReportingStrategy=ReportDFAndInfo(),
    τ_strat::TimeStepStrategy=ConstantTimeStep(),
    threading=:auto,
    m_strat::MemoryStrategy=NoMemory(),
    replica::ReplicaStrategy=NoStats(),
    post_step=(),
    maxlength= 2 * _n_walkers(v, s_strat) + 100, # padding for small walker numbers
)
    # Set up r_strat and params
    r_strat = refine_r_strat(r_strat)
    if !isnothing(laststep)
        params.laststep = laststep
    end
    if !isnothing(dτ)
        params.dτ = dτ
    end

    # Set up threading
    threading = select_threading_strategy(threading, _n_walkers(v, s_strat))
    wm = working_memory(threading, v)

    # Set up post_step
    if !(post_step isa Tuple)
        post_step = (post_step,)
    end

    # Set up replicas
    nreplicas = num_replicas(replica)
    if nreplicas > 1
        replicas = ntuple(nreplicas) do i
            ReplicaState(hamiltonian, deepcopy(v), deepcopy(wm), deepcopy(params), "_$i")
        end
    else
        replicas = (ReplicaState(hamiltonian, v, wm, params, ""),)
    end

    return QMCState(
        hamiltonian, replicas, Ref(Int(maxlength)),
        m_strat, r_strat, s_strat, τ_strat, threading, post_step, replica
    )
end

# Allow setting step and laststep from QMCState.
function Base.getproperty(state::QMCState, key::Symbol)
    if key == :step
        step = state.replicas[1].params.step
        return step
    elseif key == :laststep
        laststep = state.replicas[1].params.laststep
        return laststep
    elseif key == :maxlength
        return getfield(state, :maxlength)[]
    else
        return getfield(state, key)
    end
end
function Base.setproperty!(state::QMCState, key::Symbol, value)
    if key == :step
        for r in state.replicas
            r.params.step = value
        end
        return value
    elseif key == :laststep
        for r in state.replicas
            r.params.laststep = value
        end
        return value
    elseif key == :maxlength
        getfield(state, :maxlength)[] = value
        return value
    else
        # This will error
        return setfield!(state, key, value)
    end
end

function Base.show(io::IO, st::QMCState)
    print(io, "QMCState")
    if length(st.replicas) > 1
        print(io, " with ", length(st.replicas), " replicas")
    end
    print(io, "\n  H:    ", st.hamiltonian)
    print(io, "\n  step: ", st.step, " / ", st.laststep)
    print(io, "\n  replicas: ")
    for (i, r) in enumerate(st.replicas)
        print(io, "\n    $i: ", r)
    end
end

"""
    lomc!(ham::AbstractHamiltonian, [v]; kwargs...) -> df, state
    lomc!(state::QMCState, [df]; kwargs...) -> df, state

Linear operator Monte Carlo: Perform a projector quantum Monte Carlo simulation for
determining the lowest eigenvalue of `ham`. `v` can be a single starting vector. The default
choice is
```julia
v = DVec(starting_address(ham) => 10; style=IsStochasticInteger())
```
and triggers the integer walker FCIQMC algorithm. See [`DVec`](@ref) and
[`StochasticStyle`](@ref).

# Keyword arguments, defaults, and precedence:

* `params::FciqmcRunStrategy = RunTillLastStep(laststep = 100, dτ = 0.01, shift = diagonal_element(ham, starting_address(ham)))` - basic parameters of simulation state, see [`FciqmcRunStrategy`](@ref); is mutated
* `laststep` - can be used to override information otherwise contained in `params`
* `s_strat::ShiftStrategy = DoubleLogUpdate(targetwalkers = 100, ζ = 0.08, ξ = ζ^2/4)` - how to update the `shift`, see [`ShiftStrategy`](@ref)
* `maxlength = 2 * s_strat.targetwalkers + 100` - upper limit on the length of `v`; when reached, `lomc!` will abort
* `style = IsStochasticInteger()` - set [`StochasticStyle`](@ref) for default `v`; unused if `v` is specified.
* `post_step::NTuple{N,<:PostStepStrategy} = ()` - extract observables (e.g. [`ProjectedEnergy`](@ref)), see [`PostStepStrategy`](@ref).
* `replica::ReplicaStrategy = NoStats(1)` - run several synchronised simulation, see [`ReplicaStrategy`](@ref).
* `r_strat::ReportingStrategy = ReportDFAndInfo()` - how and when to report results, see [`ReportingStrategy`](@ref)
* `τ_strat::TimeStepStrategy = ConstantTimeStep()` - adjust time step dynamically, see [`TimeStepStrategy`](@ref)
* `m_strat::MemoryStrategy = NoMemory()` - experimental: inject memory noise, see [`MemoryStrategy`](@ref)
* `threading = :auto` - can be used to control the use of multithreading (overridden by `wm`)
  * `:auto` - use multithreading if `s_strat.targetwalkers ≥ 500`
  * `true` - use multithreading if available (set shell variable `JULIA_NUM_THREADS`!)
  * `false` - run on single thread
* `wm` - working memory; if set, it controls the use of multithreading and overrides `threading`; is mutated
* `df = DataFrame()` - when called with `AbstractHamiltonian` argument, a `DataFrame` can be passed into `lomc!` that will be pushed into.

# Return values

`lomc!` returns a named tuple with the following fields:

* `df`: a `DataFrame` with all statistics being reported.
* `state`: a `QMCState` that can be used for continuations.

# Example

```jldoctest
julia> add = BoseFS((1,2,3));

julia> hamiltonian = HubbardReal1D(add);

julia> df1, state = lomc!(hamiltonian);

julia> df2, _ = lomc!(state, df1; laststep=200); # Continuation run

julia> size(df1)
(100, 13)

julia> size(df2)
(200, 13)
```
"""
function lomc!(ham, v; df=DataFrame(), kwargs...)
    state = QMCState(ham, v; kwargs...)
    return lomc!(state, df)
end
function lomc!(ham; style=IsStochasticInteger(), kwargs...)
    v = DVec(starting_address(ham)=>10; style)
    return lomc!(ham, v; kwargs...)
end
# For continuation, you can pass a QMCState and a DataFrame
function lomc!(state::QMCState, df=DataFrame(); laststep=0)
    report = Report()
    if !iszero(laststep)
        state.laststep = laststep
    end

    # Sanity checks.
    step, laststep = state.step, state.laststep
    for replica in state.replicas
        ConsistentRNG.check_crng_independence(replica.v)
        @assert replica.params.step == step
        @assert replica.params.laststep == laststep
    end

    while step < laststep
        step += 1
        report!(state.r_strat, step, report, :steps, step)
        # This loop could be threaded if num_threads() == num_replicas? MPIData would need
        # to be aware of the replica's id and use that in communication.
        success = true
        for replica in state.replicas
            success &= advance!(report, state, replica)
        end
        replica_names, replica_values = replica_stats(state.replica, state.replicas)
        report!(state.r_strat, step, report, replica_names, replica_values)
        report_after_step(state.r_strat, step, report, state)
        ensure_correct_lengths(report)
        !success && break
    end

    # Put report into DataFrame and merge with `df`. We are assuming the previous `df` is
    # compatible, which should be the case if the run is an actual continuation. Maybe the
    # DataFrames should be merged in a more permissive manner?
    result_df = finalize_report!(state.r_strat, report)
    if !isempty(df)
        return (; df=vcat(df, result_df), state)
    else
        return (; df=result_df, state)
    end
end

"""
    advance!(report::Report, state::QMCState, replica::ReplicaState)

Advance the `replica` by one step. The `state` is used only to access the various strategies
involved. Steps, stats, and computed quantities are written to the `report`.

Returns `true` if the step was successful and calculation should proceed, `false` when
it should terminate.
"""
function advance!(
    report, state::QMCState, replica::ReplicaState{T}
) where {T}
    @unpack hamiltonian, m_strat, r_strat, s_strat, τ_strat, threading = state
    @unpack v, w, pnorm, params, id = replica
    @unpack step, shiftMode, shift, dτ = params
    step += 1

    # Step
    step_stat_names, raw_step_stats = fciqmc_step!(
        threading, w, hamiltonian, v, shift, dτ
    )
    shift_noise = apply_memory_noise!(w, v, shift, dτ, pnorm, m_strat)
    v, w, step_stat_values = sort_into_targets!(v, w, raw_step_stats)
    v, update_dvec_stats = update_dvec!(v)

    # Stats
    tnorm = walkernumber(v)
    len = length(v)

    # Updates
    shift, shiftMode, pnorm, proceed = update_shift(
        s_strat, shift, shiftMode, tnorm, pnorm, dτ, step, nothing, v, w
    )
    dτ = update_dτ(τ_strat, dτ, tnorm)

    @pack! params = step, shiftMode, shift, dτ
    @pack! replica = v, w, pnorm, params

    # Note: post_step must be called after packing the values.
    post_step_stats = post_step(state.post_step, replica)

    # Reporting
    report!(
        r_strat, step, report,
        (dτ, shift, shiftMode, len, norm=tnorm), id,
    )
    report!(r_strat, step, report, step_stat_names, step_stat_values, id)
    report!(r_strat, step, report, update_dvec_stats, id)
    report!(r_strat, step, report, (;shift_noise), id)
    report!(state.r_strat, step, report, post_step_stats, id)

    if len == 0
        if length(state.replicas) > 1
            @error "population in replica$(replica.id) is dead. Aborting."
        else
            @error "population is dead. Aborting."
        end
        return false
    end
    if len > state.maxlength[]
        if length(state.replicas) > 1
            @error "`maxlength` reached in replica$(replica.id). Aborting."
        else
            @error "`maxlength` reached. Aborting."
        end
        return false
    end
    return proceed # Bool
end


text:
# common.jl

##############################
# Booleans in the type domain
##############################

# We introduce these types to compute with booleans in a type context.
# They are not exported, but they are necessary when an external package wishes
# to extend the embeddings and promotions of spaces in this package.
const True = Val{true}
const False = Val{false}

# Simple boolean operations on the new types
(|)(::Type{True}, ::Type{True}) = True
(|)(::Type{True}, ::Type{False}) = True
(|)(::Type{False}, ::Type{True}) = True
(|)(::Type{False}, ::Type{False}) = False

# Return True if one of the arguments is True
one_of(::Type{True}) = True
one_of(::Type{False}) = False
one_of(a::Type{Val{A}}, b::Type{Val{B}}) where {A,B} = |(one_of(a),one_of(b))
one_of(a::Type{Val{A}}, b::Type{Val{B}}, c::Type{Val{C}}, d...) where {A,B,C} = one_of(a, one_of(b, c, d...))

# Convert the boolean type to a boolean value
result(::Type{True}) = true
result(::Type{False}) = false

##############################
# Promotion helper functions
##############################

"Return True if S promotes to T, i.e., if promote_type(S,T) == T."
promotes_to(S, T) = _promotes_to(S, T, promote_type(S,T))
_promotes_to(::Type{S}, ::Type{T}, ::Type{T}) where {S,T} = True
_promotes_to(::Type{S}, ::Type{T}, ::Type{U}) where {S,T,U} = False

#######################
# Composite structures
#######################

"""
Some types have composite structure, e.g. product domains, a union of domains.
These types contain a list of domains.

It is often undesirable to use `getindex` to access the elements of the composite
type. For this reason we introduce the `elements` functions. Composite types
can implement `elements` and provide a generic way to access their components.

`elements(t)`: returns the elements making up the composite type `t`

`element(t, i)`: return the `i`-th element of the composite type `t`

`numelements(t)`: return the number of elements of the composite type `t`
"""
elements() = nothing

"""
Return the i-th element of a composite structure.

See also: `elements`.
"""
element(t, i) = elements(t)[i]
# By default, we index elements(o)

"""
Return the number of elements of a composite structure.

See also: `elements`.
"""# By default, we return length(elements(t))
numelements(t) = length(elements(t))

###############
# Type factory
###############

"""
A `TypeFactory{T}` is a convenience type to simplify construction of a type.

Having `t = TypeFactory{T}` overrides `getindex` such that `t[a]` invokes `T(a)`.

For example:
```
v = TypeFactory{SVector}
v[0.1,0.2]
```
makes an `SVector{2,Float64}`.
"""
struct TypeFactory{T}
end

Base.getindex(v::TypeFactory{T}, args...) where {T} = T(args...)

const v = TypeFactory{SVector}()

###############
# Type conversion
###############
Base.convert(::Type{SVector}, ::Type{NTuple{N,T}}) where {N,T} = SVector{N,T}


text:
module FundStabABM

include("types.jl")
include("params.jl")
include("functions.jl")
include("runmodel.jl")

end


text:
using MCMCChains: MCMCChains
using CmdStan

const noncentered_schools_stan_model = """
    data {
        int<lower=0> J;
        real y[J];
        real<lower=0> sigma[J];
    }
    parameters {
        real mu;
        real<lower=0> tau;
        real eta[J];
    }
    transformed parameters {
        real theta[J];
        for (j in 1:J)
            theta[j] = mu + tau * eta[j];
    }
    model {
        mu ~ normal(0, 5);
        tau ~ cauchy(0, 5);
        eta ~ normal(0, 1);
        y ~ normal(theta, sigma);
    }
    generated quantities {
        vector[J] log_lik;
        vector[J] y_hat;
        for (j in 1:J) {
            log_lik[j] = normal_lpdf(y[j] | theta[j], sigma[j]);
            y_hat[j] = normal_rng(theta[j], sigma[j]);
        }
    }
"""

function makechains(
    names, ndraws, nchains, domains=[Float64 for _ in names]; seed=42, internal_names=[]
)
    rng = MersenneTwister(seed)
    nvars = length(names)
    var_vals = [rand(rng, domain, ndraws, nchains) for domain in domains]
    vals = permutedims(cat(var_vals...; dims=3), (1, 3, 2))
    chns = sort(MCMCChains.Chains(vals, names, Dict(:internals => internal_names)))
    return chns
end

function makechains(nvars::Int, args...; kwargs...)
    names = ["var$(i)" for i in 1:nvars]
    return makechains(names, args...; kwargs...)
end

function cmdstan_noncentered_schools(data, draws, chains; tmpdir=joinpath(pwd(), "tmp"))
    model_name = "school8"
    stan_model = Stanmodel(;
        name=model_name,
        model=noncentered_schools_stan_model,
        nchains=chains,
        num_warmup=draws,
        num_samples=draws,
        output_format=:mcmcchains,
        tmpdir=tmpdir,
    )
    rc, chns, cnames = stan(stan_model, data; summary=false)
    outfiles = ["$(tmpdir)/$(model_name)_samples_$(i).csv" for i in 1:chains]
    return (model=stan_model, files=outfiles, chains=chns)
end

function test_chains_data(chns, idata, group, names=names(chns); coords=Dict(), dims=Dict())
    ndraws, nvars, nchains = size(chns)
    @test idata isa InferenceData
    @test group in propertynames(idata)
    ds = getproperty(idata, group)
    sizes = dimsizes(ds)
    @test length(sizes) == 2 + length(coords)
    vars = vardict(ds)
    for name in names
        # `vars`, the value in ArviZ/Python is always String,
        # while `names` is String or Symbol which depends on version of MCMCChains
        name = string(name)

        @test name in keys(vars)
        dim = get(dims, name, [])
        s = (x -> length(get(coords, x, []))).(dim)
        @test size(vars[name]) == (nchains, ndraws, s...)
    end
    @test attributes(ds)["inference_library"] == "MCMCChains"
    return nothing
end

@testset "from_mcmcchains" begin
    @testset "posterior" begin
        nvars, nchains, ndraws = 2, 4, 20
        chns = makechains(nvars, ndraws, nchains)
        idata = from_mcmcchains(chns)
        test_chains_data(chns, idata, :posterior, names(chns))
    end

    @testset "prior" begin
        nvars, nchains, ndraws = 2, 4, 20
        chns = makechains(nvars, ndraws, nchains)
        idata = from_mcmcchains(; prior=chns)
        test_chains_data(chns, idata, :prior, names(chns))
    end

    @testset "posterior + prior" begin
        nvars, nchains, ndraws = 2, 4, 20
        chns = makechains(nvars, ndraws, nchains)
        idata = from_mcmcchains(chns; prior=chns)
        test_chains_data(chns, idata, :posterior, names(chns))
        test_chains_data(chns, idata, :prior, names(chns))
    end

    @testset "coords/dim" begin
        names = ["a[1]", "a[2]", "b[1]", "b[2]"]
        coords = Dict("ai" => 1:2, "bi" => ["b1", "b2"])
        dims = Dict("a" => ["ai"], "b" => ["bi"])
        nchains, ndraws = 4, 20
        chns = makechains(names, ndraws, nchains)
        idata = from_mcmcchains(chns; coords=coords, dims=dims)
        test_chains_data(chns, idata, :posterior, ["a", "b"]; coords=coords, dims=dims)
        vardims = dimdict(idata.posterior)
        @test vardims["a"] == ("chain", "draw", "ai")
        @test vardims["b"] == ("chain", "draw", "bi")
    end

    @testset "multivariate" begin
        names = ["a[1][1]", "a.2.2", "a[2,1]", "a[1, 2]"]
        coords = Dict("ai" => 1:2, "aj" => ["aj1", "aj2"])
        dims = Dict("a" => ["ai", "aj"])
        nchains, ndraws = 4, 20
        chns = makechains(names, ndraws, nchains)

        # String or Symbol, which depends on MCMCChains version
        ET = Base.promote_typeof(chns.name_map.parameters...)

        idata = from_mcmcchains(chns; coords=coords, dims=dims)
        test_chains_data(chns, idata, :posterior, ["a"]; coords=coords, dims=dims)
        arr = vardict(idata.posterior)["a"]
        @test arr[:, :, 1, 1] == permutedims(chns.value[:, ET(names[1]), :], [2, 1])
        @test arr[:, :, 2, 2] == permutedims(chns.value[:, ET(names[2]), :], [2, 1])
        @test arr[:, :, 2, 1] == permutedims(chns.value[:, ET(names[3]), :], [2, 1])
        @test arr[:, :, 1, 2] == permutedims(chns.value[:, ET(names[4]), :], [2, 1])
    end

    @testset "specify eltypes" begin
        # https://github.com/arviz-devs/ArviZ.jl/issues/125
        nchains, ndraws = 4, 20
        names = ["x", "y", "z"]
        domains = [Float64, (0, 1), 1:3]
        post = makechains(names, ndraws, nchains, domains)
        prior = makechains(names, ndraws, nchains, domains)
        post_pred = makechains(["d"], ndraws, nchains, [(0, 1)])
        idata = from_mcmcchains(
            post;
            prior=prior,
            posterior_predictive=post_pred,
            eltypes=Dict("y" => Bool, "z" => Int, "d" => Bool),
        )
        test_chains_data(post, idata, :posterior)
        test_chains_data(prior, idata, :prior)
        test_chains_data(post_pred, idata, :posterior_predictive)
        @test eltype(idata.posterior[:y].values) <: Bool
        @test eltype(idata.posterior[:z].values) <: Int64
        @test eltype(idata.prior[:y].values) <: Bool
        @test eltype(idata.prior[:z].values) <: Int64
        @test eltype(idata.posterior_predictive[:d].values) <: Bool

        idata2 = from_mcmcchains(
            post;
            prior=prior,
            posterior_predictive=["y"],
            eltypes=Dict("y" => Bool, "z" => Int),
        )
        test_chains_data(post, idata2, :posterior, ["x", "z"])
        test_chains_data(prior, idata2, :prior)
        @test eltype(idata2.posterior[:z].values) <: Int64
        @test eltype(idata2.prior[:y].values) <: Bool
        @test eltype(idata2.prior[:z].values) <: Int64
        @test eltype(idata2.posterior_predictive[:y].values) <: Bool
    end

    @testset "complete" begin
        nchains, ndraws = 4, 20
        nobs = 10
        posterior = prior = ["a[1]", "a[2]"]
        posterior_predictive = prior_predictive = ["ahat[1]", "ahat[2]"]
        observed_data = Dict("xhat" => 1:nobs)
        constant_data = Dict("x" => (1:nobs) .+ nobs)
        predictions_constant_data = Dict("z" => (1:nobs) .+ nobs)
        predictions = "zhat"
        sample_stats = ["stat"]
        log_likelihood = "log_lik"
        post_names = [
            posterior
            posterior_predictive
            sample_stats
            predictions
            log_likelihood
        ]
        prior_names = [prior; prior_predictive; sample_stats]
        chns = makechains(post_names, ndraws, nchains; internal_names=["stat"])
        chns2 = makechains(prior_names, ndraws, nchains; internal_names=["stat"])
        idata = from_mcmcchains(
            chns;
            posterior_predictive="ahat",
            predictions="zhat",
            prior=chns2,
            prior_predictive="ahat",
            observed_data=observed_data,
            constant_data=constant_data,
            predictions_constant_data=predictions_constant_data,
            log_likelihood="log_lik",
        )
        for group in (
            :posterior,
            :prior,
            :posterior_predictive,
            :prior_predictive,
            :predictions,
            :log_likelihood,
            :observed_data,
            :constant_data,
            :predictions_constant_data,
            :sample_stats,
            :sample_stats_prior,
        )
            @test group in propertynames(idata)
        end
        @test length(dimdict(idata.posterior)) == 4
        @test "a" ∈ keys(dimdict(idata.posterior))
        @test length(dimdict(idata.posterior_predictive)) == 4
        @test "ahat" ∈ keys(dimdict(idata.posterior_predictive))
        @test length(dimdict(idata.prior)) == 4
        @test "a" ∈ keys(dimdict(idata.prior))
        @test length(dimdict(idata.prior_predictive)) == 4
        @test "ahat" ∈ keys(dimdict(idata.prior_predictive))
        @test length(dimdict(idata.sample_stats)) == 3
        @test "stat" ∈ keys(dimdict(idata.sample_stats))
        @test length(dimdict(idata.predictions)) == 3
        @test "zhat" ∈ keys(dimdict(idata.predictions))
        @test length(dimdict(idata.log_likelihood)) == 3
        @test "log_lik" ∈ keys(dimdict(idata.log_likelihood))
        @test length(dimdict(idata.sample_stats_prior)) == 3
        @test "stat" ∈ keys(dimdict(idata.sample_stats_prior))
        @test length(dimdict(idata.observed_data)) == 2
        @test "xhat" ∈ keys(dimdict(idata.observed_data))
        @test length(dimdict(idata.constant_data)) == 2
        @test "x" ∈ keys(dimdict(idata.constant_data))
        @test length(dimdict(idata.predictions_constant_data)) == 2
        @test "z" ∈ keys(dimdict(idata.predictions_constant_data))
    end

    @testset "missing -> NaN" begin
        rng = MersenneTwister(42)
        nvars, nchains, ndraws = 2, 4, 20
        vals = randn(rng, ndraws, nvars, nchains)
        vals = Array{Union{Float64,Missing},3}(vals)
        vals[1, 1, 1] = missing
        names = ["var$(i)" for i in 1:nvars]
        chns = MCMCChains.Chains(vals, names)
        @test Missing <: eltype(chns.value)
        idata = from_mcmcchains(chns)
        vdict = vardict(idata.posterior)
        @test eltype(vdict["var1"]) <: Real
        @test isnan(vdict["var1"][1, 1, 1])
    end

    @testset "info -> attributes" begin
        nvars, nchains, ndraws = 2, 4, 20
        chns = makechains(nvars, ndraws, nchains)
        idata = from_mcmcchains(chns; library="MyLib")
        attrs = idata.posterior.attrs
        @test attrs isa Dict
        @test attrs["inference_library"] == "MyLib"
    end
end

@testset "convert_to_dataset(::MCMCChains.Chains)" begin
    nvars, nchains, ndraws = 2, 4, 20
    chns = makechains(nvars, ndraws, nchains)
    ds = ArviZ.convert_to_dataset(chns; library="MyLib")
    @test ds isa ArviZ.Dataset
    attrs = ds.attrs
    @test "inference_library" ∈ keys(attrs)
    @test attrs["inference_library"] == "MyLib"
end

@testset "convert_to_inference_data(::MCMCChains.Chains)" begin
    nvars, nchains, ndraws = 2, 4, 20
    chns = makechains(nvars, ndraws, nchains)
    idata = convert_to_inference_data(chns; group=:posterior)
    @test idata isa InferenceData
    @test :posterior in propertynames(idata)
    idata = convert_to_inference_data(chns; group=:prior)
    @test idata isa InferenceData
    @test :prior in propertynames(idata)
end

@testset "test MCMCChains readme example" begin
    # Define the experiment
    n_iter = 500
    n_name = 3
    n_chain = 2

    # experiment results
    val = randn(n_iter, n_name, n_chain) .+ [1, 2, 3]'
    val = hcat(val, rand(1:2, n_iter, 1, n_chain))

    # construct a Chains object
    chn = MCMCChains.Chains(val)  # According to version, this may introduce String or Symbol name

    @test ArviZ.summary(chn) !== nothing
end

VERSION > v"1.0" && @testset "from_cmdstan" begin
    data = noncentered_schools_data()
    mktempdir() do path
        output = cmdstan_noncentered_schools(data, 500, 4; tmpdir=path)
        posterior_predictive = prior_predictive = ["y_hat"]
        log_likelihood = "log_lik"
        coords = Dict("school" => 1:8)
        dims = Dict(
            "theta" => ["school"],
            "y" => ["school"],
            "log_lik" => ["school"],
            "y_hat" => ["school"],
            "eta" => ["school"],
        )
        idata1 = from_cmdstan(
            output.chains;
            posterior_predictive=posterior_predictive,
            log_likelihood=log_likelihood,
            prior=output.chains,
            prior_predictive=prior_predictive,
            coords=coords,
            dims=dims,
        )
        idata2 = from_cmdstan(
            output.files;
            posterior_predictive=posterior_predictive,
            log_likelihood=log_likelihood,
            prior=output.files,
            prior_predictive=prior_predictive,
            coords=coords,
            dims=dims,
        )
        @testset "idata.$(group)" for group in Symbol.(idata2._groups)
            ds1 = getproperty(idata1, group)
            ds2 = getproperty(idata2, group)
            for f in (vardict, dimsizes)
                d1 = f(ds1)
                d2 = f(ds2)
                for k in keys(d1)
                    @test k in keys(d2)
                    @test d1[k] ≈ d2[k]
                end
            end
        end
    end
end


text:
push!(LOAD_PATH,"../src/")

using Documenter

using Genie, Genie.App, Genie.AppServer, Genie.Assets
using Genie.Cache, Genie.Commands, Genie.Configuration, Genie.Cookies
using Genie.Encryption, Genie.FileTemplates, Genie.Flax, Genie.Generator
using Genie.Inflector, Genie.Input, Genie.Plugins
using Genie.Renderer, Genie.Requests, Genie.Responses, Genie.Router
using Genie.Sessions, Genie.Tester, Genie.Toolbox, Genie.Util, Genie.WebChannels

push!(LOAD_PATH,  "../../src",
                  "../../src/cache_adapters",
                  "../../src/session_adapters")

makedocs(sitename = "Genie - The Highly Productive Julia Web Framework", format = Documenter.HTML(prettyurls = false))


text:
immutable LogLoss <: Loss
end

@inline function loss_value(loss::LogLoss, p::Real, y::Real)
    return -log(ifelse(y == 1.0, p, 1 - p))
end

Base.show(io::IO, loss::LogLoss) = print(io, "Log-Loss")

@inline function loss_derivative(loss::LogLoss, p::Real, y::Real)
    return p - y
end


text:
using AmbientForcing
using OrdinaryDiffEq, Distributions

function rober(du,u,p,t)
  y₁,y₂,y₃ = u
  du[1] = -0.04 * y₁ + 1e4 * y₂ * y₃
  du[2] =  0.04 * y₁ - 1e4 * y₂ * y₃ - 3e7 * y₂^2
  du[3] =  y₁ + y₂ + y₃ - 1
  nothing
end

M = [1. 0  0
     0  1. 0
     0  0  0];

ode_rober = ODEFunction(rober, mass_matrix = M)

u0 = [1.0,0.0,0.0];

g_rober = constraint_equations(ode_rober)
g_rober(u0)

Frand = random_force(ode_rober, [0.0,1], Uniform)
z_new = ambient_forcing(ode_rober, u0, 2.0, Frand)

g_rober(z_new)

h = [0, 1 ,0]
z_new = ambient_forcing(ode_rober, u0, 2.0, h)

g_rober(z_new)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


text:
testOnAllBackends("Binary Operators") do af
	println("\tarithmetic")
	afArr = array(af, [[1, 2] [3, 4]])

	println("\t./")
	result = afArr ./ 2
	@test host(result) == [[0, 1] [1, 2]]

	println("\taddAssign!")
	addAssign!(result, 1)
	@test host(result) == [[1, 2] [2, 3]]

	println("\t.*=")
	result .*= 2
	@test host(result) == [[2, 4] [4, 6]]

	println("\tsubAssign!")
	subAssign!(result, 1.5f0)
	@test host(result) == [[0.5f0, 2.5f0] [2.5f0, 4.5f0]]

	println("\t.\\=")
	result .\= 2.0
	@test host(result) == [[4.0, 0.8] [0.8, 0.4444444444444444]]

	afArr = array(af, [[1, 2] [3, 4]])

	println("\t./")
	result = afArr ./ array(af, [[2, 2] [2, 2]])
	@test host(result) == [[0, 1] [1, 2]]

	println("\taddAssign!")
	addAssign!(result, array(af, [[1, 1] [1, 1]]))
	@test host(result) == [[1, 2] [2, 3]]

	println("\tmulAssign!")
	mulAssign!(result, array(af, [[2, 2] [2, 2]]))
	@test host(result) == [[2, 4] [4, 6]]

	println("\t.-=")
	result .-= array(af, [[1.5f0, 1.5f0] [1.5f0, 1.5f0]])
	@test host(result) == [[0.5f0, 2.5f0] [2.5f0, 4.5f0]]

	println("\t.\=")
	result .\= array(af, [[2.0, 2.0] [2.0, 2.0]])
	@test host(result) == [[4.0, 0.8] [0.8, 0.4444444444444444]]

	println("\tlogic")
	afArr = array(af, [[1, 2] [3, 4]])

	println("\t.<")
	result = afArr .< array(af, [[2, 2] [2, 2]])
	@test host(result) == [[true, false] [false, false]]

	println("\t.<=")
	result = afArr .<= array(af, [[2, 2] [2, 2]])
	@test host(result) == [[true, true] [false, false]]

	println("\t.==")
	result = afArr .== array(af, [[2, 2] [2, 2]])
	@test host(result) == [[false, true] [false, false]]

	println("\t.!=")
	result = afArr .!= array(af, [[2, 2] [2, 2]])
	@test host(result) == [[true, false] [true, true]]

	println("\t.>")
	result = afArr .> array(af, [[2, 2] [2, 2]])
	@test host(result) == [[false, false] [true, true]]

	println("\t.>=")
	result = afArr .>= 2
	@test host(result) == [[false, true] [true, true]]

	afArr = array(af, [[0, 1] [2, 0]])

	println("\tand")
	result = and(afArr, array(af, [[true, false] [true, false]]))
	@test host(result) == [[false, false] [true, false]]

	println("\tor")
	result = or(afArr, array(af, [[true, false] [true, false]]))
	@test host(result) == [[true, true] [true, false]]

	println("\tmaxOf, minOf")
	afArr1 = array(af, [[1.0f0, 2.0f0] [4.0f0, 3.0f0]])
	afArr2 = array(af, [[5.0f0, 1.0f0] [1.0f0, 5.0f0]])

	@test host(maxOf(afArr1, afArr2)) == [[5.0f0, 2.0f0] [4.0f0, 5.0f0]]
	@test host(minOf(afArr1, 2)) == [[1.0f0, 2.0f0] [2.0f0, 2.0f0]]
	@test host(maxOf(2.0, afArr2)) == [[5.0, 2.0] [2.0, 5.0]]
end

max_stars_repo_name: unbornchikken/julia-ml-proto

max_issues_repo_name: unbornchikken/julia-ml-proto

max_forks_repo_name: unbornchikken/julia-ml-proto


text:
# includes required for parallel processes

using JuMP
using LightGraphs
using LightGraphsFlows
# using Gurobi
using GLPK

include("constants.jl")
include("transport.jl")
include("basic.jl")
include("intermediate.jl")


text:
"""
    f(lx,ly,...) type methods deprecated; see analysis.
"""
fourierenergy(lx::T,ly::T,nx::Int,ny::Int,u::Array{DNSField{T},1}) where T = energyspectrum.(Ref(Domain(lx,ly,nx,ny)),u)
fourierenergy(lx::T,ly::T,nx::Int,ny::Int,u::Array{DSSField{T},1}) where T = energyspectrum.(Ref(Domain(lx,ly,nx,ny)),u)
fourierenergy(lx::T,ly::T,nx::Int,ny::Int,Λ::Int,u::Array{GSSField{T},1}) where T = energyspectrum.(Ref(Domain(lx,ly,nx,ny)),u)

"""
    Runtime rank computations deprecated; see analysis.
"""
function rankis(cumulant::Array{ComplexF64,4},nx::Int,ny::Int,Λ::Int)
    twopoint = reshape(cumulant,(2*ny-1)*(nx-Λ),(2*ny-1)*(nx-Λ))
    D = eigvals(twopoint)
    return ((2*ny-1)*(nx-Λ) - length(D[D .< 1e-6]),D)
end

function rankis(cumulant::Array{ComplexF64,3},nx::Int,ny::Int)
    temp::Array{ComplexF64,4} = zeros(ComplexF64,2*ny-1,nx-1,2*ny-1,nx-1)
    modalevs::Array{Float64,2} = zeros(Float64,2*ny-1,nx-1)
    modalranks::Array{Int64,1} = zeros(Int64,nx-1)
    for m1=1:nx-1
        modalevs[:,m1] = eigvals(cumulant[:,:,m1])
        modalranks[m1] = length(modalevs[:,m1][modalevs[:,m1] .> 1e-6])
        for n1=-(ny-1):ny-1
            for n2=-(ny-1):ny-1

                 temp[n2+ny,m1,n1+ny,m1] = cumulant[n2+ny,n1+ny,m1]
            end
        end
    end
    twopoint = reshape(temp,(2*ny-1)*(nx-1),(2*ny-1)*(nx-1))
    D = eigvals(twopoint)
    return ((2*ny-1)*(nx-1) - length(D[D .< 1e-6]),D,modalranks,modalevs)
end

"""
    Custom noise functions deprecated; standard Gaussian used in solve.
"""
function dist!(dW::DNSField{T},W,dt,u,p,t,rng) where T<:AbstractFloat
    nx,ny,F = p.nx,p.ny,p.F
    @inbounds for m=1:nx-1
        @inbounds for n=-ny+1:ny-1
            ϕ = rand(Uniform(0,2π))
            dW[n+ny,m+1] = abs(sqrt(dt))*F[n+ny,m+1]*exp(im*ϕ)
        end
    end
    nothing
end

function dist!(dW::GSSField{T},W,dt,u,p,t,rng) where T<:AbstractFloat
    nx,ny,Λ,F = p.nx,p.ny,p.Λ,p.F
    @inbounds for m=1:Λ
        @inbounds for n=-ny+1:ny-1
            ϕ = rand(Uniform(0,2π))
            dW.x[1][n+ny,m+1] = abs(sqrt(dt))*F.x[1][n+ny,m+1]*exp(im*ϕ)
        end
    end
    dW.x[2] .= zero(Complex{T})
    nothing
end

# bridge!(dW,W,W0,Wh,q,h,u,p,t,rng) = dW .= W0 .+ h .* (Wh .- W0)
bridge!(dW,W,W0,Wh,q,h,u,p,t,rng) = @. dW = W0 + h * (Wh - W0)
noise!(t0,W0,Z0=nothing;kwargs...) = NoiseProcess(t0,W0,Z0,dist!,bridge!;kwargs...)


text:
struct ShamirSecretSharing
    prime::BigInt
end

ShamirSecretSharing() = ShamirSecretSharing(BigInt(2)^127 - 1)

"""
n is the number of shares.
k is the threshold.

"""
function Base.split(scheme::ShamirSecretSharing, secret::BigInt, n::Int64, k::Int64)
    prime = scheme.prime
    polynomial = [secret]
    for i in 2:k
        push!(polynomial, rand(RandomDevice(), 0:(prime-1)))
    end

    points = []
    for i in 1:n
        push!(points, (i, eval_poly(polynomial, i, prime)))
    end
    return points
end

function Base.join(scheme::ShamirSecretSharing, points)
    prime = scheme.prime
    xs, ys = zip(points...)
    return lagrange_interpolate(0, [xs...], [ys...], prime)
end


text:
using LinearAlgebra

M = [
    182.57507334 136.93130833  91.28754047  45.64377071
    136.93130757 102.69848317  68.46565663  34.23282867
     91.2875396   68.46565636  45.6437717   22.82188609
     45.64377017  34.23282846  22.82188603  11.41094314]
println(pinv(M))

M = [171.24513133 128.4338515   85.6225691   42.81128498
128.43385082  96.32539036  64.21692798  32.10846431
 85.62256831  64.21692774  42.81128587  21.40564315
 42.81128449  32.10846412  21.4056431   10.70282166]
println(pinv(M))