text:
using Revise
using Statistics
using LaTeXStrings
using PyPlot
using LinearAlgebra
using Infiltrator
using Bem2d

"""
    discretized_arc(θstart, θend, radius, n_pts)

Generate regularly spaced eleemnts along an curved arc.
"""
function discretized_arc(θstart, θend, radius, n_pts)
    # Create geometry of discretized arc
    θrange = collect(LinRange(θstart, θend, n_pts + 1))
    x = @. radius * cos(θrange)
    y = @. radius * sin(θrange)
    x1 = x[1:1:end-1]
    x2 = x[2:1:end]
    y1 = y[1:1:end-1]
    y2 = y[2:1:end]
    return x1, y1, x2, y2
end

"""
    circle_subplot(nrows, ncols, plotidx, x, y, mat, npts, R, theta0, title_string)

Plot field (displacement, stress) within a circular disk and style
"""
function circle_subplot(nrows, ncols, plotidx, x, y, mat, npts, R, theta0, title_string)
    fontsize = 20
    contour_levels = 100
    contour_color = "white"
    contour_linewidth = 0.5
    color_scale = 1

    subplot(nrows, ncols, plotidx)
    contourf(reshape(x, npts, npts), reshape(y, npts, npts), reshape(mat, npts, npts), levels=contour_levels)
    cbar = colorbar(fraction=0.05, pad=0.05, extend = "both")
    cbar.ax.tick_params(labelsize=fontsize)
    contour(reshape(x, npts, npts), reshape(y, npts, npts), reshape(mat, npts, npts), levels=contour_levels, colors=contour_color, linewidths=contour_linewidth)
    xlabel("x (m)", fontsize=fontsize)
    ylabel("y (m)", fontsize=fontsize)
    title(title_string, fontsize=fontsize)

    # Draw entire circle and region of applied tractions
    x1, y1, x2, y2 = discretized_arc(deg2rad(-180), deg2rad(180), R, 360)
    plot([x1, x2], [y1, y2], "-k", linewidth=2)
    x1, y1, x2, y2 = discretized_arc(-theta0, theta0, R, 50)
    plot([x1, x2], [y1, y2], "-r", linewidth=2)
    x1, y1, x2, y2 = discretized_arc(-theta0+deg2rad(180), theta0+deg2rad(180), R, 50)
    plot([x1, x2], [y1, y2], "-r", linewidth=2)
    gca().set_aspect("equal")
    gca().tick_params(labelsize=fontsize)
end

"""
    calcbrazil(p, x, y, R, theta0)

Calculate stresses predicted from Hondros solution to Brazil test.
"""
function calcbrazil(p, x, y, R, theta0)
    #! Solution from Hondros (1959) as summarized by Wei and Chau 2013
    r = @. sqrt(x^2 + y^2)
    theta = @. atan(y, x)

    # Analytic stresses in cylindrical coordinates
    Srr = zeros(length(x))
    Sthetatheta = zeros(length(x))
    Srtheta = zeros(length(x))
    Srrconstterm = 2.0 * theta0 * -p / deg2rad(180)
    Sthetathetaconstterm = 2.0 * theta0 * -p / deg2rad(180)
    leadingterm = 2.0 * -p / deg2rad(180)
    mmax = 1000 # Max number of terms in Hondros series
    for m in 1:mmax
        Srr += @. (r/R)^(2*m-2) * (1-(1-1/m)*(r/R)^2) * sin(2*m*theta0) * cos(2*m*theta)
        Sthetatheta += @. (r/R)^(2*m-2) * (1-(1+1/m)*(r/R)^2) * sin(2*m*theta0) * cos(2*m*theta)
        Srtheta += @. ((r/R)^(2*m) - (r/R)^(2*m-2)) * sin(2*m*theta0) * sin(2*m*theta)
    end
    Srr = @. Srrconstterm + leadingterm * Srr
    Sthetatheta = @. Sthetathetaconstterm - leadingterm * Sthetatheta
    Srtheta = @. leadingterm * Srtheta

    # Convert analytic cylindrical stresses to Cartesian
    Sanalytic = zeros(length(x), 3)
    for i in 1:length(x) # Project a single matrix to Cartesian coordinates
        cylindrical_stress_tensor = [Srr[i] Srtheta[i] ; Srtheta[i] Sthetatheta[i]]
        transformation_matrix = [cos(theta[i]) -sin(theta[i]) ; sin(theta[i]) cos(theta[i])]
        cartesian_stress_tensor = transformation_matrix * cylindrical_stress_tensor * transpose(transformation_matrix)
        Sanalytic[i, 1] = cartesian_stress_tensor[1, 1]
        Sanalytic[i, 2] = cartesian_stress_tensor[2, 2]
        Sanalytic[i, 3] = cartesian_stress_tensor[1, 2]
    end
    return Sanalytic
end

"""
    braziltest()

Compare the Hondros (1959 analytic soluiton) for the Brazil test
with the BEM solution for both constant and quadratic element cases.
"""
function braziltest()
    close("all")
    mu = 3e10
    nu = 0.25
    p = -1.0e5 # Applied radial pressure over arc
    theta0 = deg2rad(1.0) # Arc length over which pressure is applied
    nels = 360
    R = nels / (2 * pi)
    npts = 200
    x, y = Bem2d.obsgrid(-R, -R, R, R, npts)
    r = @. sqrt(x^2 + y^2)

    #! Analytic solution
    Sanalytic = calcbrazil(p, x, y, R, theta0)

    #! BEM solution
    els = Bem2d.Elements(Int(1e5))
    x1, y1, x2, y2 = discretized_arc(deg2rad(-180), deg2rad(180), R, nels)
    for i in 1:length(x1)
        els.x1[els.endidx + i] = x1[i]
        els.y1[els.endidx + i] = y1[i]
        els.x2[els.endidx + i] = x2[i]
        els.y2[els.endidx + i] = y2[i]
        els.name[els.endidx + i] = "circle"
    end
    Bem2d.standardize_elements!(els)
    idx = Bem2d.getidxdict(els)

    #! Apply normal tractions everywhere and convert from radial to Cartesian
    xtracC = zeros(els.endidx)
    ytracC = zeros(els.endidx)
    thetaels = @. atan(els.ycenter[1:1:els.endidx], els.xcenter[1:1:els.endidx])
    for i in 1:els.endidx # Calcuate the x and y components of the tractions
        normalTractions = [0; p] # Pressure in fault normal component only.
        xtracC[i], ytracC[i] = els.rotmat[i, :, :] * normalTractions
    end

    #! Zero out the tractions on the area without contact
    deleteidx = findall(x -> (x>theta0 && x<deg2rad(180)-theta0), thetaels)
    xtracC[deleteidx] .= 0
    ytracC[deleteidx] .= 0
    deleteidx = findall(x -> (x<-theta0 && x>-deg2rad(180)+theta0), thetaels)
    xtracC[deleteidx] .= 0
    ytracC[deleteidx] .= 0

    #! Kernels, T*: displacement to displacement, H*: displacement to traction
    @time _, _, HstarC = partialsconstdispstress(slip2dispstress, els, idx["circle"], idx["circle"], mu, nu)
    @time _, _, HstarQ = partialsquaddispstress(slip2dispstress, els, idx["circle"], idx["circle"], mu, nu)

    #! CONSTANT CASE
    #! Applied tractions -> effective displacements -> internal stresses
    DeffC = inv(HstarC) * interleave(xtracC, ytracC)
    @time _, SdispC = constdispstress(slip2dispstress, x, y, els, idx["circle"], DeffC[1:2:end], DeffC[2:2:end], mu, nu)

    #! QUADRATIC CASE
    #! Applied tractions -> effective displacements -> internal stresses
    xtracQ = zeros(3 * length(xtracC))
    ytracQ = zeros(3 * length(ytracC))
    xtracQ[1:3:end] = xtracC
    xtracQ[2:3:end] = xtracC
    xtracQ[3:3:end] = xtracC
    ytracQ[1:3:end] = ytracC
    ytracQ[2:3:end] = ytracC
    ytracQ[3:3:end] = ytracC
    DeffQ = inv(HstarQ) * interleave(xtracQ, ytracQ)
    @time _, SdispQ = quaddispstress(slip2dispstress, x, y, els, idx["circle"], quadstack(DeffQ[1:2:end]), quadstack(DeffQ[2:2:end]), mu, nu)

    #! Isolate the values inside the circle
    nanidx = findall(x -> x > R, r)
    Sanalytic[nanidx, :] .= NaN
    SdispC[nanidx, :] .= NaN
    SdispQ[nanidx, :] .= NaN
    SresidualC = @. SdispC - Sanalytic
    SresidualQ = @. SdispQ - Sanalytic

    #! Summary figure
    figure(figsize=(30,20))
    fontsize = 20
    nrows = 4
    ncols = 6

    #! Constant dislocation case
    # Applied tractions
    subplot(nrows, ncols, 1)
    plot(rad2deg.(thetaels), xtracC, ".r")
    plot(rad2deg.(thetaels), ytracC, "+b")
    title("applied tractions", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Effective displacements
    subplot(nrows, ncols, 2)
    plot(rad2deg.(thetaels), DeffC[1:2:end], ".r")
    plot(rad2deg.(thetaels), DeffC[2:2:end], "+b")
    title("Effective displacements", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Analytic solutions
    circle_subplot(nrows, ncols, 7, x, y, Sanalytic[:, 1], npts, R, theta0, "Sxx (analytic)")
    circle_subplot(nrows, ncols, 8, x, y, Sanalytic[:, 2], npts, R, theta0, "Syy (analytic)")
    circle_subplot(nrows, ncols, 9, x, y, Sanalytic[:, 3], npts, R, theta0, "Sxy (analytic)")

    # BEM solutions
    circle_subplot(nrows, ncols, 13, x, y, SdispC[:, 1], npts, R, theta0, "Sxx (DDM)")
    circle_subplot(nrows, ncols, 14, x, y, SdispC[:, 2], npts, R, theta0, "Syy (DDM)")
    circle_subplot(nrows, ncols, 15, x, y, SdispC[:, 3], npts, R, theta0, "Sxy (DDM)")
    circle_subplot(nrows, ncols, 19, x, y, SresidualC[:, 1], npts, R, theta0, "Sxx (residual)")
    circle_subplot(nrows, ncols, 20, x, y, SresidualC[:, 2], npts, R, theta0, "Syy (residual)")
    circle_subplot(nrows, ncols, 21, x, y, SresidualC[:, 3], npts, R, theta0, "Sxy (residual)")

    #! Quadratic case
    # Applied tractions
    subplot(nrows, ncols, 4)
    plot(xtracQ, ".r")
    plot(ytracQ, "+b")
    title("applied tractions", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Effective displacements
    subplot(nrows, ncols, 5)
    plot(DeffQ[1:2:end], ".r")
    plot(DeffQ[2:2:end], "+b")
    title("Effective displacements", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Analytic solutions
    circle_subplot(nrows, ncols, 10, x, y, Sanalytic[:, 1], npts, R, theta0, "Sxx (analytic)")
    circle_subplot(nrows, ncols, 11, x, y, Sanalytic[:, 2], npts, R, theta0, "Syy (analytic)")
    circle_subplot(nrows, ncols, 12, x, y, Sanalytic[:, 3], npts, R, theta0, "Sxy (analytic)")

    # BEM solutions
    circle_subplot(nrows, ncols, 16, x, y, SdispQ[:, 1], npts, R, theta0, "Sxx (DDM)")
    circle_subplot(nrows, ncols, 17, x, y, SdispQ[:, 2], npts, R, theta0, "Syy (DDM)")
    circle_subplot(nrows, ncols, 18, x, y, SdispQ[:, 3], npts, R, theta0, "Sxy (DDM)")
    circle_subplot(nrows, ncols, 22, x, y, SresidualQ[:, 1], npts, R, theta0, "Sxx (residual)")
    circle_subplot(nrows, ncols, 23, x, y, SresidualQ[:, 2], npts, R, theta0, "Syy (residual)")
    circle_subplot(nrows, ncols, 24, x, y, SresidualQ[:, 3], npts, R, theta0, "Sxy (residual)")
    tight_layout()
    show()

    #! 9-panel plot for quadratic only
    figure(figsize=(15,15))
    nrows = 3
    ncols = 3
    # Analytic solutions
    circle_subplot(nrows, ncols, 1, x, y, Sanalytic[:, 1], npts, R, theta0, L"\sigma_{xx} \; \mathrm{(analytic)}")
    circle_subplot(nrows, ncols, 2, x, y, Sanalytic[:, 2], npts, R, theta0, L"\sigma_{yy} \; \mathrm{(analytic)}")
    circle_subplot(nrows, ncols, 3, x, y, Sanalytic[:, 3], npts, R, theta0, L"\sigma_{xy} \; \mathrm{(analytic)}")

    # BEM solutions
    circle_subplot(nrows, ncols, 4, x, y, SdispQ[:, 1], npts, R, theta0, L"\sigma_{xx} \; \mathrm{(BEM)}")
    circle_subplot(nrows, ncols, 5, x, y, SdispQ[:, 2], npts, R, theta0, L"\sigma_{yy} \; \mathrm{(BEM)}")
    circle_subplot(nrows, ncols, 6, x, y, SdispQ[:, 3], npts, R, theta0, L"\sigma_{xy} \; \mathrm{(BEM)}")
    circle_subplot(nrows, ncols, 7, x, y, SresidualQ[:, 1], npts, R, theta0, L"\sigma_{xx} \; \mathrm{(residual)}")
    circle_subplot(nrows, ncols, 8, x, y, SresidualQ[:, 2], npts, R, theta0, L"\sigma_{yy} \; \mathrm{(residual)}")
    circle_subplot(nrows, ncols, 9, x, y, SresidualQ[:, 3], npts, R, theta0, L"\sigma_{xy} \; \mathrm{(residual)}")
    tight_layout()
    show()
end
braziltest()


text:
include( "../src/HouseholdMappingFRS_HBAI.jl")
create_data()

ext: jl

max_stars_repo_name: grahamstark/ScottishTaxBenefitModel

max_issues_repo_name: grahamstark/ScottishTaxBenefitModel

max_forks_repo_name: grahamstark/ScottishTaxBenefitModel


text:
using Flux
using WindMixing
using JLD2
using FileIO
using OceanParameterizations
using OrdinaryDiffEq
using Random
using GalacticOptim
using LinearAlgebra

BLAS.set_num_threads(1)

# Training data
# train_files = [
#     "wind_-1e-3_heating_-4e-8",
#     "wind_-5e-4_cooling_4e-8",
#     ]

train_files = [
    "-1e-3"
]

𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

PATH = pwd()

OUTPUT_PATH = joinpath(PATH, "training_output")
OUTPUT_PATH = "D:\\University Matters\\MIT\\CLiMA Project\\OceanParameterizations.jl\\training_output"

EXTRACTED_OUTPUT_PATH = joinpath(PATH, "extracted_training_output")

# FILE_PATH = joinpath(OUTPUT_PATH, "NDE_training_mpp_3sim_-1e-3_-8e-4_-5e-4_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_1e-2_rate_2e-4.jld2")
# @assert !isfile(FILE_PATH)

# FILE_PATH_uw = joinpath(PATH, "extracted_training_output", "uw_NN_training_1sim_-1e-3_extracted.jld2")
# FILE_PATH_vw = joinpath(PATH, "extracted_training_output", "vw_NN_training_1sim_-1e-3_extracted.jld2")
# FILE_PATH_wT = joinpath(PATH, "extracted_training_output", "wT_NN_training_1sim_-1e-3_extracted.jld2")

# uw_file = jldopen(FILE_PATH_uw, "r")
# vw_file = jldopen(FILE_PATH_vw, "r")
# wT_file = jldopen(FILE_PATH_wT, "r")

# uw_NN = uw_file["neural_network"]
# vw_NN = vw_file["neural_network"]
# wT_NN = wT_file["neural_network"]

N_inputs = 96
hidden_units = 400
N_outputs = 31

weights, re = Flux.destructure(Chain(Dense(N_inputs, hidden_units, relu), Dense(hidden_units, N_outputs)))

uw_NN = re(weights ./ 1f5)
vw_NN = re(weights ./ 1f5)
wT_NN = re(weights ./ 1f5)

task_id = parse(Int,ARGS[1]) + 1
num_tasks = parse(Int,ARGS[2])
task_id = 1
FILE_NAME = ["NDE_training_mpp_2sim_windcooling_MS_windheating_SS_diffusivity_1e-1_Ri_1e-1_divide1f5_gradient_smallNN_scale_5e-3_rate_1e-4",
             "NDE_training_mpp_2sim_windcooling_MS_windheating_SS_diffusivity_1e-1_Ri_1e-1_divide1f5_gradient_smallNN_scale_1e-2_rate_1e-4",
             "NDE_training_mpp_2sim_windcooling_MS_windheating_SS_diffusivity_1e-1_Ri_1e-1_divide1f5_gradient_smallNN_scale_1.5e-2_rate_1e-4",
             "NDE_training_mpp_2sim_windcooling_MS_windheating_SS_diffusivity_1e-1_Ri_1e-1_divide1f5_gradient_smallNN_scale_2e-2_rate_1e-4",
              ][task_id]

FILE_PATH = joinpath(OUTPUT_PATH, "$(FILE_NAME).jld2")
@assert !isfile(FILE_PATH)

EXTRACTED_FILE_NAME = "$(FILE_NAME)_extracted"
EXTRACTED_FILE_PATH = joinpath(EXTRACTED_OUTPUT_PATH, "$EXTRACTED_FILE_NAME.jld2")

# FILE_NAME_NN = ["NDE_training_mpp_5sim_-1e-3_-9e-4_-8e-4_-7e-4_-5e-4_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_5e-3_rate_1e-4_extracted.jld2",
#                 "NDE_training_mpp_5sim_-1e-3_-9e-4_-8e-4_-7e-4_-5e-4_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_5e-3_rate_2e-4_extracted.jld2",
#                 "NDE_training_mpp_5sim_-1e-3_-9e-4_-8e-4_-7e-4_-5e-4_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_1e-2_rate_1e-4_extracted.jld2",
#                 "NDE_training_mpp_5sim_-1e-3_-9e-4_-8e-4_-7e-4_-5e-4_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_1e-2_rate_2e-4_extracted.jld2"
#                 ][task_id]

# FILE_NAME_NN = "NDE_training_mpp_9sim_windcooling_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_5e-3_rate_1e-4_extracted.jld2"
# FILE_PATH_NN = joinpath(EXTRACTED_OUTPUT_PATH, FILE_NAME_NN)

# @assert isfile(FILE_PATH_NN)
# file = jldopen(FILE_PATH_NN, "r")

# uw_NN = file["neural_network/uw"]
# vw_NN = file["neural_network/vw"]
# wT_NN = file["neural_network/wT"]

# close(file)

gradient_scaling = [5f-3, 1f-2, 1.5f-2, 2f-2][task_id]
train_parameters = Dict("ν₀" => 1f-4, "ν₋" => 0.1f0, "Riᶜ" => 0.25f0, "ΔRi" => 1f-1, "Pr" => 1f0, "κ" => 10f0,
                        "modified_pacanowski_philander" => true, "convective_adjustment" => false,
                        "smooth_profile" => false, "smooth_NN" => false, "smooth_Ri" => false, "train_gradient" => true,
                        "zero_weights" => true, "unscaled" => false, "gradient_scaling" => gradient_scaling)

# train_epochs = [1]
# train_tranges = [1:9:1153]
# train_iterations = [600]
# # train_optimizers = [[[ADAM(1e-4)]], [[ADAM(2e-4)]], [[ADAM(1e-4)]], [[ADAM(2e-4)]], [[ADAM(1e-4)]], [[ADAM(2e-4)]], [[ADAM(1e-4)]], [[ADAM(2e-4)]]][task_id]
# train_optimizers = [[ADAM(1e-4)]]

train_epochs = [1]
train_tranges = [1:35:200]
train_iterations = [3]
train_optimizers = [[ADAM(2e-4)]]

# train_tranges = [1:10:100, 1:10:200, 1:20:500, 1:30:700, 1:30:800, 1:30:900, 1:35:1153]
# train_epochs = [1 for i in 1:length(train_tranges)]
# train_iterations = [50, 50, 100, 30, 20, 50, 150]
# train_optimizers = [[ADAM(0.1), ADAM(0.01)], [ADAM(0.01)], [ADAM(0.01)], [ADAM(0.01)], [ADAM(0.01)], [ADAM(0.01)], [ADAM(0.01), ADAM(0.001), ADAM(5e-4), ADAM(2e-4)]]

# train_tranges = [1:10:100, 1:10:200, 1:20:500, 1:20:800, 1:35:1153]
# train_epochs = [1 for i in 1:length(train_tranges)]
# train_iterations = [30, 30, 50, 30, 200]
# train_optimizers = [[[ADAM(0.01)] for i in 1:6]; [[ADAM(0.01), ADAM(1e-3)]]]
# # train_optimizers = [[ADAM(1e-5)] for i in 1:6]

timestepper = Rosenbrock23()

# train_optimizers = [[ADAM(2e-4), ADAM(1e-4), ADAM(5e-5), RMSProp(1e-4)]]
# train_optimizers=[[ADAM(5e-4)]]

function train(FILE_PATH, train_files, train_epochs, train_tranges, train_parameters, train_optimizers, train_iterations, uw_NN, vw_NN, wT_NN, 𝒟train, timestepper, unscaled)
    write_metadata_NDE_training(FILE_PATH, train_files, train_epochs, train_tranges, train_parameters, train_optimizers, uw_NN, vw_NN, wT_NN)
    if unscaled
        for i in 1:length(train_epochs)
            @info "iteration $i/$(length(train_epochs)), time range $(train_tranges[i])"
            # uw_NN, vw_NN, wT_NN = train_NDE_convective_adjustment(uw_NN, vw_NN, wT_NN, 𝒟train, train_tranges[i], timestepper, train_optimizers[i], train_epochs[i], FILE_PATH, 1, 1, 10f0, 5)
            if train_parameters["modified_pacanowski_philander"]
                uw_NN, vw_NN, wT_NN = train_NDE_unscaled(uw_NN, vw_NN, wT_NN, 𝒟train, train_tranges[i], timestepper, train_optimizers[i], train_epochs[i], FILE_PATH, i, n_simulations=length(train_files), maxiters=train_iterations[i],
                modified_pacanowski_philander=train_parameters["modified_pacanowski_philander"], convective_adjustment=train_parameters["convective_adjustment"],
                ν₀=train_parameters["ν₀"], ν₋=train_parameters["ν₋"], ΔRi=train_parameters["ΔRi"], Riᶜ=train_parameters["Riᶜ"],
                κ=train_parameters["κ"],
                smooth_profile=train_parameters["smooth_profile"], smooth_NN=train_parameters["smooth_NN"], smooth_Ri=train_parameters["smooth_Ri"], train_gradient=train_parameters["train_gradient"],
                gradient_scaling=train_parameters["gradient_scaling"])
            else
                uw_NN, vw_NN, wT_NN = train_NDE_unscaled(uw_NN, vw_NN, wT_NN, 𝒟train, train_tranges[i], timestepper, train_optimizers[i], train_epochs[i], FILE_PATH, i, n_simulations=length(train_files), maxiters=train_iterations[i],
                modified_pacanowski_philander=train_parameters["modified_pacanowski_philander"], convective_adjustment=train_parameters["convective_adjustment"],
                κ=train_parameters["κ"],
                smooth_profile=train_parameters["smooth_profile"], smooth_NN=train_parameters["smooth_NN"], smooth_Ri=train_parameters["smooth_Ri"], train_gradient=train_parameters["train_gradient"],
                gradient_scaling=train_parameters["gradient_scaling"])
            end
        end
    else
        for i in 1:length(train_epochs)
            @info "iteration $i/$(length(train_epochs)), time range $(train_tranges[i])"
            # uw_NN, vw_NN, wT_NN = train_NDE_convective_adjustment(uw_NN, vw_NN, wT_NN, 𝒟train, train_tranges[i], timestepper, train_optimizers[i], train_epochs[i], FILE_PATH, 1, 1, 10f0, 5)
            if train_parameters["modified_pacanowski_philander"]
                uw_NN, vw_NN, wT_NN = train_NDE(uw_NN, vw_NN, wT_NN, 𝒟train, train_tranges[i], timestepper, train_optimizers[i], train_epochs[i], FILE_PATH, i, n_simulations=length(train_files), maxiters=train_iterations[i],
                modified_pacanowski_philander=train_parameters["modified_pacanowski_philander"], convective_adjustment=train_parameters["convective_adjustment"],
                ν₀=train_parameters["ν₀"], ν₋=train_parameters["ν₋"], ΔRi=train_parameters["ΔRi"], Riᶜ=train_parameters["Riᶜ"],
                κ=train_parameters["κ"],
                smooth_profile=train_parameters["smooth_profile"], smooth_NN=train_parameters["smooth_NN"], smooth_Ri=train_parameters["smooth_Ri"], train_gradient=train_parameters["train_gradient"],
                zero_weights = train_parameters["zero_weights"],
                gradient_scaling=train_parameters["gradient_scaling"])
            else
                uw_NN, vw_NN, wT_NN = train_NDE(uw_NN, vw_NN, wT_NN, 𝒟train, train_tranges[i], timestepper, train_optimizers[i], train_epochs[i], FILE_PATH, i, n_simulations=length(train_files), maxiters=train_iterations[i],
                modified_pacanowski_philander=train_parameters["modified_pacanowski_philander"], convective_adjustment=train_parameters["convective_adjustment"],
                κ=train_parameters["κ"],
                smooth_profile=train_parameters["smooth_profile"], smooth_NN=train_parameters["smooth_NN"], smooth_Ri=train_parameters["smooth_Ri"], train_gradient=train_parameters["train_gradient"],
                zero_weights = train_parameters["zero_weights"],
                gradient_scaling=train_parameters["gradient_scaling"])
            end
        end
    end
    return uw_NN, vw_NN, wT_NN
end

# uw_NN_res, vw_NN_res, wT_NN_res = train(FILE_PATH, train_files, train_epochs, train_tranges, train_parameters, train_optimizers, train_iterations, uw_NN, vw_NN, wT_NN, 𝒟train, timestepper, train_parameters["unscaled"])

uw_NN_res, vw_NN_res, wT_NN_res =  train(FILE_PATH, train_files, train_epochs, train_tranges, train_parameters, train_optimizers, train_iterations, uw_NN, vw_NN, wT_NN, 𝒟train, timestepper, train_parameters["unscaled"])

extract_NN(FILE_PATH, EXTRACTED_FILE_PATH, "NDE")

test_files = train_files

animate_training_results(test_files, FILE_NAME, trange=1:1:1153)

OCEANANIGANS_OUTPUT_DIR = joinpath(pwd(), "NDE_output_oceananigans", FILE_NAME)
animate_training_results_oceananigans(test_files, 60, FILE_NAME, OCEANANIGANS_OUTPUT_DIR)


text:
########################################################################################
# Integrated Brownian Motion
########################################################################################
"""
    ibm(d::Integer, q::Integer, elType=typeof(1.0))

Generate the discrete dynamics for a q-IBM model. INCLUDES AUTOMATIC PRECONDITIONING!
"""
function ibm(d::Integer, q::Integer, elType=typeof(1.0))
    # Make A
    A_breve = zeros(elType, q + 1, q + 1)
    @simd ivdep for j in 1:q+1
        @simd ivdep for i in 1:j
            @inbounds A_breve[i, j] = binomial(q - i + 1, q - j + 1)
        end
    end
    A = kron(I(d), A_breve)
    @assert istriu(A)
    # A = UpperTriangular(A)

    # Make Q
    Q_breve = zeros(elType, q + 1, q + 1)
    @fastmath _transdiff_ibm_element(row::Int, col::Int) =
        one(elType) / (2 * q + 1 - row - col)
    @simd ivdep for col in 0:q
        @simd ivdep for row in 0:q
            val = _transdiff_ibm_element(row, col)
            @inbounds Q_breve[1+row, 1+col] = val
        end
    end
    QL_breve = cholesky!(Q_breve).L
    QL = kron(I(d), QL_breve)
    Q = SRMatrix(QL)

    return A, Q
end

"""Same as above, but without the automatic preconditioning"""
function vanilla_ibm(d::Integer, q::Integer)
    @fastmath function A!(A::AbstractMatrix, h::Real)
        # Assumes that A comes from a previous computation => zeros and one-diag
        val = one(h)
        for i in 1:q
            val = val * h / i
            for k in 0:d-1
                for j in 1:q+1-i
                    @inbounds A[j+k*(q+1), j+k*(q+1)+i] = val
                end
            end
        end
    end

    @fastmath function _transdiff_ibm_element(row::Int, col::Int, h::Real)
        idx = 2 * q + 1 - row - col
        fact_rw = factorial(q - row)
        fact_cl = factorial(q - col)
        return h^idx / (idx * fact_rw * fact_cl)
    end
    @fastmath function Q!(Q::AbstractMatrix, h::Real, σ²::Real=1.0)
        val = one(h)
        @simd for col in 0:q
            @simd for row in col:q
                val = _transdiff_ibm_element(row, col, h) * σ²
                @simd for i in 0:d-1
                    @inbounds Q[1+col+i*(q+1), 1+row+i*(q+1)] = val
                    @inbounds Q[1+row+i*(q+1), 1+col+i*(q+1)] = val
                end
            end
        end
    end

    return A!, Q!
end


text:
using StructuresKit

#test case
#8ZS2.25x070, no rotational spring, no moment gradient, Lm>>Lcrd
#sharp corners

CorZ = 1

t = 0.070

#out to out dimensions
ho = 8.0
b = 2.25
d = 0.930
θc = 50.0

E = 29500.0
μ = 0.30
G = E/(2*(1+μ))
f1 = 1.0
f2 = -1.0
M1 = 1.0
M2 = 1.0
Lm = 100.0
kϕ = 0.0

Sf = 9.4251/4.0

CurvatureSign = -1.0

Mcrd = AISIS10016.app23331(CorZ, t, ho, b, d, θc, E, μ, G, f1, f2, M1, M2, CurvatureSign, Lm, kϕ, Sf)

#AISI equations underpredict for a Z?
@test Mcrd ≈ 121.86 atol=0.09*Mcrd


text:
using Parameters

@with_kw struct sp_params
    # force parameters
    ϵ::Float64  = 10            # interaction distance (percentage)
    sC::Float64 = 1.0           # cohesion strength coefficient
    sA::Float64 = 1.0           # adhesion strength coefficient
    G::Float64  = 1.0           # force of gravity

    # time stepping
    Tf::Float64 = 10              # final time
    Δt::Float64 = 0.1            # maximum time step

    # geometry
    circfile::String = "initial_figures.circ"    # file name for particle geometry
    wallfile::String = "step.wall"               # file name for wall geometry
    matID = Dict(1 => :clay)                     # IDs for particles

    # don't touch these (used internally)
    Nt::Int = Int(ceil(Tf/Δt))   # number of steps to take
    # need term for 'largest step size'
end


text:
import Dates

import ..AstroDates: Date, DateTime, Time,
    year, month, day,
    hour, minute, second, millisecond,
    fractionofday, yearmonthday, dayofyear

"""
    timescale(ep)

Return the time scale of epoch `ep`.

# Example

```jldoctest
julia> ep = TTEpoch(2000, 1, 1)
2000-01-01T00:00:00.000 TT

julia> timescale(ep)
TT
```
"""
timescale(ep::Epoch) = ep.scale

"""
    DateTime(ep::Epoch)

Convert the epoch `ep` to an `AstroDates.DateTime`.
"""
function DateTime(ep::Epoch)
    if !isfinite(ep.fraction)
        if ep.fraction < 0
            return DateTime(AstroDates.MIN_EPOCH, AstroDates.H00)
        else
            return DateTime(AstroDates.MAX_EPOCH, Time(23, 59, 59.999))
        end
    end

    sec = ep.second + Int64(43200)
    time = sec % Int64(86400)
    if time < 0
        time += Int64(86400)
    end
    date = Int((sec - time) ÷ Int64(86400))

    date_comp = Date(AstroDates.J2000_EPOCH, date)
    time_comp = Time(time, ep.fraction)

    leap = getleap(ep)
    hr = hour(time_comp)
    mn = minute(time_comp)
    leap = ifelse(hr == 23 && mn == 59 && abs(leap) == 1.0, leap, 0.0)
    if !iszero(leap)
        h = hour(time_comp)
        m = minute(time_comp)
        s = second(Float64, time_comp) + leap
        time_comp = Time(h, m, s)
    end

    DateTime(date_comp, time_comp)
end

"""
    year(ep::Epoch)

Get the year of the epoch `ep`.
"""
year(ep::Epoch) = year(DateTime(ep))

"""
    month(ep::Epoch)

Get the month of the epoch `ep`.
"""
month(ep::Epoch) = month(DateTime(ep))

"""
    day(ep::Epoch)

Get the day of the epoch `ep`.
"""
day(ep::Epoch) = day(DateTime(ep))

"""
    yearmonthday(ep::Epoch)

Get the year, month, and day of the epoch `ep` as a tuple.
"""
yearmonthday(ep::Epoch) = yearmonthday(DateTime(ep))

"""
    dayofyear(ep::Epoch)

Get the day of the year of the epoch `ep`.
"""
dayofyear(ep::Epoch) = dayofyear(DateTime(ep))

"""
    hour(ep::Epoch)

Get the hour of the epoch `ep`.
"""
hour(ep::Epoch) = hour(DateTime(ep))

"""
    minute(ep::Epoch)

Get the minute of the epoch `ep`.
"""
minute(ep::Epoch) = minute(DateTime(ep))

"""
    second(type, ep::Epoch)

Get the second of the epoch `ep` as a `type`.
"""
second(typ, ep::Epoch) = second(typ, DateTime(ep))

"""
    second(ep::Epoch) -> Int

Get the second of the epoch `ep` as an `Int`.
"""
second(ep::Epoch) = second(Int, DateTime(ep))

"""
    millisecond(ep::Epoch)

Get the number of milliseconds of the epoch `ep`.
"""
millisecond(ep::Epoch) = millisecond(DateTime(ep))

"""
    Time(ep::Epoch)

Get the `Time` of the epoch `ep`.
"""
Time(ep::Epoch) = Time(DateTime(ep))

"""
    Date(ep::Epoch)

Get the `Date` of the epoch `ep`.
"""
Date(ep::Epoch) = Date(DateTime(ep))

"""
    fractionofday(ep::Epoch)

Get the time of the day of the epoch `ep` as a fraction.
"""
fractionofday(ep::Epoch) = fractionofday(Time(ep))

"""
    Dates.DateTime(ep::Epoch)

Convert the epoch `ep` to a `Dates.DateTime`.
"""
Dates.DateTime(ep::Epoch) = Dates.DateTime(DateTime(ep))

"""
    Dates.Date(ep::Epoch)

Convert the date of epoch `ep` to a `Dates.Date`.
"""
Dates.Date(ep::Epoch) = Dates.Date(Date(ep))

"""
    Dates.Time(ep::Epoch)

Convert the time of epoch `ep` to a `Dates.Time`.
"""
Dates.Time(ep::Epoch) = Dates.Time(Time(ep))


text:
            # P is the precision used with the typed occurance
            #
type ArbFloat{P}  <: Real     # field and struct names from arb.h
  mid_exp ::Int               #           fmpz
  mid_size::UInt              #           mp_size_t
  mid_d1  ::Int               #           mantissa_struct
  mid_d2  ::Int               #
  rad_exp ::Int               #           fmpz
  rad_man ::UInt              #
end

precision{P}(x::ArbFloat{P}) = P
precision{P}(::Type{ArbFloat{P}}) = P
# allow inquiring the precision of the module: precision(ArbFloats)
precision(::Type{Type{Val{:ArbFloats}}}) = precision(ArbFloat)
precision(m::Module) = precision(Type{Val{Symbol(m)}})

# precision is significand precision, significand_bits(FloatNN) + 1, for the hidden bit
typealias ArbFloat16  ArbFloat{ 11}  # read  2 ? 3 or fewer decimal digits to write the same digits ( 16bit Float)
typealias ArbFloat32  ArbFloat{ 24}  # read  6 ? 7 or fewer decimal digits to write the same digits ( 32bit Float)
typealias ArbFloat64  ArbFloat{ 53}  # read 15 ?15 or fewer decimal digits to write the same digits ( 64bit Float)
typealias ArbFloat128 ArbFloat{113}  # read 33 ?34 or fewer decimal digits to write the same digits (128bit Float)
typealias ArbFloat256 ArbFloat{237}  # read 71 ?71 or fewer decimal digits to write the same digits (256bit Float)
typealias ArbFloat512 ArbFloat{496}  # read 148?149 or fewer decimal digits to write the same digits (512bit Float)

# get and set working precision for ArbFloat

const ArbFloatPrecision = [116,]
precision(::Type{ArbFloat}) = ArbFloatPrecision[1]

function setprecision(::Type{ArbFloat}, x::Int)
    priorprecision = precision(ArbFloat)
    x = max(11, abs(x))
    x > 4095 && warn("ArbFloats are designed to work best at precisions < 4096 bits")
    ArbFloatPrecision[1] = x
    priorprecision
end

# a type specific hash function helps the type to 'just work'
const hash_arbfloat_lo = (UInt === UInt64) ? 0x37e642589da3416a : 0x5d46a6b4
const hash_0_arbfloat_lo = hash(zero(UInt), hash_arbfloat_lo)
# two values of the same precision
#    with identical midpoint significands and identical radial exponents hash equal
# they are the same value, one is less accurate yet centered about the other
hash{P}(z::ArbFloat{P}, h::UInt) =
    hash(reinterpret(UInt,z.mid_d1)$z.mid_exp,
         (h $ hash(reinterpret(UInt,z.mid_d2)$(~reinterpret(UInt,P)), hash_arbfloat_lo)
            $ hash_0_arbfloat_lo))

# adapted from Nemo
function (==){P}(x::ArbFloat{P}, y::ArbFloat{P})
    return Bool(ccall(@libarb(arb_eq), Cint, (Ptr{ArbFloat{P}}, Ptr{ArbFloat{P}}), &x, &y))
end
function (!=){P}(x::ArbFloat{P}, y::ArbFloat{P})
    return Bool(ccall(@libarb(arb_ne), Cint, (Ptr{ArbFloat{P}}, Ptr{ArbFloat{P}}), &x, &y))
end

function clearArbFloat{P}(x::ArbFloat{P})
     ccall(@libarb(arb_clear), Void, (Ptr{ArbFloat{P}},), &x)
end

function initializer{P}(::Type{ArbFloat{P}})
    z = ArbFloat{P}(0,0,0,0,0,0)
    ccall(@libarb(arb_init), Void, (Ptr{ArbFloat{P}},), &z)
    finalizer(z, clearArbFloat)
    z
end

function midpoint{P}(x::ArbFloat{P})
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_get_mid_arb), Void, (Ptr{ArbFloat}, Ptr{ArbFloat}), &z, &x)
    z
end

function radius{P}(x::ArbFloat{P})
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_get_rad_arb), Void, (Ptr{ArbFloat}, Ptr{ArbFloat}), &z, &x)
    z
end

function upperbound{P}(x::ArbFloat{P})
    a = ArfFloat{P}(0,0,0,0)
    ccall(@libarb(arf_init), Void, (Ptr{ArfFloat{P}},), &a)
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_get_ubound_arf), Void, (Ptr{ArfFloat}, Ptr{ArbFloat}, Int), &a, &x, P)
    ccall(@libarb(arb_set_arf), Void, (Ptr{ArbFloat}, Ptr{ArfFloat}), &z, &a)
    ccall(@libarb(arf_clear), Void, (Ptr{ArfFloat{P}},), &a)
    z
end

function lowerbound{P}(x::ArbFloat{P})
    a = ArfFloat{P}(0,0,0,0)
    ccall(@libarb(arf_init), Void, (Ptr{ArfFloat{P}},), &a)
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_get_lbound_arf), Void, (Ptr{ArfFloat}, Ptr{ArbFloat}, Int), &a, &x, P)
    ccall(@libarb(arb_set_arf), Void, (Ptr{ArbFloat}, Ptr{ArfFloat}), &z, &a)
    ccall(@libarb(arf_clear), Void, (Ptr{ArfFloat{P}},), &a)
    z
end

bounds{P}(x::ArbFloat{P}) = ( lowerbound(x), upperbound(x) )

function minmax{P}(x::ArbFloat{P}, y::ArbFloat{P})
   x < y ? (x,y) : (y,x)
end

function relativeError{P}(x::ArbFloat{P})
    z = P
    ccall(@libarb(arb_rel_error_bits), Void, (Ptr{Int}, Ptr{ArbFloat}), &z, &x)
    z
end

function relativeAccuracy{P}(x::ArbFloat{P})
    z = P
    ccall(@libarb(arb_rel_accuracy_bits), Void, (Ptr{Int}, Ptr{ArbFloat}), &z, &x)
    z
end

function midpointPrecision{P}(x::ArbFloat{P})
    z = P
    ccall(@libarb(arb_bits), Void, (Ptr{Int}, Ptr{ArbFloat}), &z, &x)
    z
end

function trimmedAccuracy{P}(x::ArbFloat{P})
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_trim), Void, (Ptr{ArbFloat}, Ptr{ArbFloat}), &z, &x)
    z
end