make_orbits.jl

@everywhere push!(LOAD_PATH, "/home/stagnerl/")
import InterpolatingFunctions; @everywhere using InterpolatingFunctions
import GuidingCenterOrbits; @everywhere using GuidingCenterOrbits
import HDF5; @everywhere using HDF5
import Clustering; @everywhere using Clustering
import Iterators; @everywhere using Iterators
import StaticArrays; @everywhere using StaticArrays
import FastKmeans; @everywhere using FastKmeans

function make_clustered_distribution_file(M,wall,nkk)
    nenergy = 100
    npitch = 100
    nr = 100
    energy = linspace(5.0,80.0,nenergy)
    pitch = linspace(-1.0,1.0,npitch)
    R = linspace(1.73,2.3,nr)
    println("Starting Orbit Calculation")
    @time begin
        orbs = @parallel (vcat) for k=1:nr
            begin
                orbs_k = Array(Orbit{Float64},(nenergy,npitch))
                cc = EPRCoordinate(M, 0.0, 0.0, R[k])
                for i=1:nenergy, j=1:npitch
                    c = EPRCoordinate(energy[i],pitch[j],cc.r,cc.z,cc.amu,cc.q)
                    o = get_orbit(M, c, nstep=6000,tmax=600.0)
                    if o.class == :incomplete || hits_wall(o,wall)
                        orbs_k[i,j] = Orbit(o.coordinate)
                    else
                        orbs_k[i,j] = o
                    end
                end
                orbs_k
            end
        end
    end
#    get_rz = function g(r)
#        c = EPRCoordinate(M, 0.0, 0.0, r)
#        return (c.r,c.z)
#    end
#    @time rz = pmap(get_rz,R)
#    RR = getindex.(rz,1)
#    ZZ = getindex.(rz,2)
#
#    calc_orbit = function f(eprz)
#        c = EPRCoordinate(eprz...)
#        o = get_orbit(M, c, nstep=6000,tmax=600.0)
#        if o.class == :incomplete || hits_wall(o,wall)
#            return Orbit(o.coordinate)
#        else
#            return o
#        end
#    end
#    @time orbs = pmap(calc_orbit, product(energy,pitch,RR,ZZ))

    oclasses = [:co_passing, :ctr_passing, :potato, :trapped, :stagnation]
    orbit_counts = Dict{Symbol,Int}(o=>0 for o in oclasses)
    orbit_parts = Dict{Symbol,Int}(o=>0 for o in oclasses)
    println("Starting Clustering")
    @time begin
        ec = Float64[]; sizehint!(ec,Int(1e5))
        pc = Float64[]; sizehint!(pc,Int(1e5))
        rc = Float64[]; sizehint!(rc,Int(1e5))
        inds = Int64[]; sizehint!(inds,Int(1e5))
        oclass = Symbol[]
        npart = 0
        norbits = 0
        for (i,o) in enumerate(orbs)
            o.class == :incomplete && continue
            orbit_counts[o.class] += 1
            orbit_parts[o.class] = orbit_parts[o.class] + length(o.path.r)
            npart = npart + length(o.path.r)
            norbits = norbits + 1
            push!(ec,o.coordinate.energy)
            push!(pc,o.coordinate.pitch)
            push!(rc,o.coordinate.r)
            push!(oclass,o.class)
            push!(inds,i)
        end
        ocounts = [orbit_counts[o] for o in oclasses]
        oi = sortperm(ocounts)
        oclasses = oclasses[oi]

        for nk in nkk
            println(nk)
            nclusters = 0
            for oc in oclasses
                nnk = max(Int(round(nk*orbit_counts[oc]/norbits)),1)
                println(oc,": ",orbit_counts[oc]," ",nnk)
                if (nclusters + nnk) > nk
                    nnk = nk - nclusters
                end

                npart = orbit_parts[oc]
                wo = oclass .== oc
                if nnk > 1
                    opoints = [[eec/75,ppc/2,rrc/(2.3-1.73)] for (eec,ppc,rrc) in zip(ec[wo],pc[wo],rc[wo])]
                    @time ass, centers = fetch(@spawn fastkmeans(opoints,nnk))
                   # @time k = fetch(@spawn kmeans([ec[wo]./75 pc[wo]./2 rc[wo]./(2.3-1.73)]', nnk))
                   # ass = k.assignments
                   # centers = k.centers
                else
                    ass = fill(1,sum(wo))
                    centers = [[mean(ec[wo]),mean(pc[wo]),mean(rc[wo])]./[75.0,2.0,(2.3-1.73)]]
                    #centers = reshape([mean(ec[wo]),mean(pc[wo]),mean(rc[wo])]./[75.0,2.0,(2.3-1.73)], (3,1))
                end

                ee = Float64[]; sizehint!(ee,npart)
                pp = Float64[]; sizehint!(pp,npart)
                zz = Float64[]; sizehint!(zz,npart)
                rr = Float64[]; sizehint!(rr,npart)
                weight = Float64[]; sizehint!(weight,npart)
                kass = Int16[]; sizehint!(kass,npart)
                class = Int64[]; sizehint!(class,npart)
                for (i,ind) in enumerate(inds[wo])
                    o = orbs[ind]
                    np = length(o.path.r)
                    append!(ee,fill(o.coordinate.energy,np))
                    #pitch = fetch(@spawn get_pitch(M, o.coordinate, o.path))
                    append!(pp,-o.path.pitch)
                    append!(rr,o.path.r*100)
                    append!(zz,o.path.z*100)
                    append!(kass,fill(Int16(ass[i]),np))
                    append!(class,fill(i,np))
                    append!(weight,o.path.dt.*(1e19/sum(o.path.dt)))
                end
                odir = "/cscratch/lstagner/orbits_"*@sprintf("%03d",nk)
                isdir(odir) || mkdir(odir)
                @parallel for i=1:nnk
                    sfile = odir*"/orbit_"*@sprintf("%03d",i+nclusters)*"_"*@sprintf("%03d",nk)*".h5"
                    w = kass .== i
                    ww = ass .== i
                    np = sum(w)
                    no = sum(ww)
                    ocl = Array(Int16,np)
                    for (j,c) in enumerate(unique(class[w]))
                        cw = c .== class[w]
                        ocl[cw] = j
                    end
                    h5open(sfile,"w") do file
                        file["orbit_type"] = string(oc)
                        file["coordinate_center"] = Array(centers[i]).*[75.0,2.0,(2.3-1.73)]
                        file["energy_c", "shuffle",(), "chunk", (no),"compress", 4] = ec[wo][ww]
                        file["pitch_c", "shuffle",(), "chunk", (no),"compress", 4] = pc[wo][ww]
                        file["r_c", "shuffle",(), "chunk", (no),"compress", 4] = rc[wo][ww]
                        file["data_source"] = "calculated from geqdsk"
                        file["type"] = 2
                        file["nclass"] = length(unique(ocl))
                        file["nparticle"] = np
                        file["time"] = 0.790
                        file["energy", "shuffle",(), "chunk", (np),"compress", 4] = ee[w]
                        file["pitch","shuffle",(), "chunk", (np),"compress", 4] = pp[w]
                        file["r","shuffle",(), "chunk", (np),"compress", 4] = rr[w]
                        file["z","shuffle",(), "chunk", (np),"compress", 4] = zz[w]
                        file["class","shuffle",(), "chunk", (np),"compress", 4] = ocl
                        file["weight","shuffle",(), "chunk", (np),"compress", 4] = weight[w]
                   end
                end
                nclusters = nclusters + nnk
            end
        end
    end
    return nothing
end

@everywhere begin
gfile = "g159243.00791.h5"
g = h5open(gfile);
fpol = Float64.(read(g["fpol"]))
pressure = Float64.(read(g["pres"]))
psirz = Float64.(read(g["psirz"]))
psiaxis = Float64(read(g["ssimag"]))
psiwall = Float64(read(g["ssibry"]))
psi_domain = (psiaxis, psiwall)

B0 = Float64(read(g["bcentr"]))
J0 = Float64(read(g["cpasma"]))
sigma = sign(B0*J0)

rzero = Float64(read(g["rzero"]))
fpol0 = B0*rzero
psi = [linspace(psiaxis,psiwall,length(fpol))...]

r = Float64.(read(g["r"]))
nr = length(r)
r_domain = extrema(r)
z = Float64.(read(g["z"]))
nz = length(z)
z_domain = extrema(z)

raxis = Float64.(read(g["rmaxis"]))
zaxis = Float64.(read(g["zmaxis"]));
rwall = Float64.(read(g["lim"])[1,:])
zwall = Float64.(read(g["lim"])[2,:]);
nb = read(g["nbdry"])
lcfs = Float64.(read(g["bdry"]))[:,1:nb];

wall = Polygon()
for i =1:length(rwall)
    push!(wall.vertices,(rwall[i],zwall[i]))
end

f_psi = interpolate(psirz, (r,z), BicubicSpline(), Irregular{2,Val{size(psirz)}}, Nearest)
f_fpol = interpolate(fpol, (psi,), CubicSpline(), Irregular(length(psi)), Value(fpol0))
f_pressure = interpolate(pressure, (psi,), CubicSpline(), Irregular(length(psi)), Value(0.0))
M = AxisymmetricEquilibrium(r_domain, z_domain, psi_domain, f_psi, f_fpol, f_pressure, (raxis,zaxis));
end
@time make_clustered_distribution_file(M, wall, [500])