Integrate NeutronTransport.jl into FUSE

[kevinm387]

I am working to integrate NeutronTransport.jl into FUSE to do neutron transport calculations. I have a pull request on NeutronTransport.jl to add fixed source calculations, and also have a new branch on FUSE with some simple functions that call NeutronTransport.jl, that I plan to expand to couple FUSE and NeutronTransport.jl.

Right now I am struggling for speed with some of the functions in NeutronTransport.jl. In particular, in mocsolver.jl, there are several functions with lots of nested for loops that use many allocations, such as _solve(), compute_q!() and compute_φ!(). It would be helpful to get some additional eyes on some of these functions, as I believe there are ways we could possibly speed this up.

Let me know what other info I could provide. Also, where would be a good place to put a Jupyter notebook with a test case that we could use to compare speeds?

Thanks!

[bclyons12]

If the test case is only a few lines of code, you could include those commands here. Otherwise, you could put the notebook in FUSE/playground

[kevinm387]

Ok, I created a file here.

This method also requires a nuclear data file that currently is about 380 MB. Do you guys have thoughts on where that should go? That file size can probably be reduced by a lot because it includes cross sections for lots of materials that we won't actually need (e.g. acetone, 20+ different kinds of concrete, and others).

[orso82]

where does that nuclear data file come from? can we download it openly from the web?

[bclyons12]

If it's not openly accessible, it could probably live on saturn/omega somewhere in the /fusion/ga/projects/ird/ptp if we'll only want to access on those clusters. If it should be bundled with the PTP packages, you can use Git LFS: https://docs.github.com/en/repositories/working-with-files/managing-large-files/about-git-large-file-storage

[kevinm387]

The nuclear data file was generated using OpenMC and a small package I made, mgxs_generator. So no, this particular nuclear data file can't be downloaded from the web.

For the time being I put it on Omega at /fusion/projects/ird/ptp/mclaughlink

[bclyons12]

I think the slowness is related to all the runtime typing. zero(T) and such. The way it's written, the compiler doesn't seem to know what type to expect, so it can't optimize it and it's just done at runtime. I'm assuming we need to make more clever use of multiple dispatching for the compiler to infer the types, but some simple testing I did had no effect. I'll keep looking.

[bclyons12]

Nevermind my previous comment. There was an enormous performance hit caused by xss being defined as a Vector{NeutronTransport.CrossSections}. NeutronTransport.CrossSections isn't a concrete type, since it is parametrically typed by

CrossSections{
    NGroups,elType,
    T<:Union{Vector{elType},SVector{NGroups,elType}},
    S<:Union{Matrix{elType},SMatrix{NGroups,NGroups,elType}}
}

Thus, this vector couldn't communicate what elType was to the compiler and compute_q! was runtime dispatching. I fixed it with this: 910f3e9

[bclyons12]

Here's the profiling of the full run after this change. Most of the time is still spent in comput_q!, but at least that's type-stable and not allocating anymore.

There is allocation/runtime dispatching happening in compute_φ! (note that it's red). I think it's because the TrackGenerator struct in RayTracing.jl is poorly written. It contains Vector{Track}, but Track is parametric, so it's not type stable. This is somewhat necessary as it appears that the tracks_by_uid field does in fact contain Tracks of different parametric types. Looks like it would require some rewriting to avoid this issue. Probably not worth it for now.

[bclyons12]

In NeutronTransport.jl, changing compute_q! to this has a >3x speedup at the cost of one intentional allocation:

function compute_q!(sol::MoCSolution{T}, prob::MoCProblem, fixed_sources::Matrix{T}) where {T}
    NGroups = ngroups(prob)
    NRegions = nregions(prob)
    @unpack keff, φ, q = sol
    φtmp = Vector{T}(undef, NGroups)

    @inbounds for i in 1:NRegions
        xs = getxs(prob, i)
        @unpack χ, Σt, νΣf, Σs0 = xs
        fissionable = isfissionable(xs)

        for g′ in 1:NGroups
            φtmp[g′] = φ[@region_index(i, g′)]
        end

        for g in 1:NGroups
            ig = @region_index(i, g)
            @views cΣs0 = Σs0[:, g]
            qig = dot(cΣs0, φtmp)
            if fissionable
                qig += sum(1 / keff * χ[g] * νΣf[g′] * φtmp[g′] for g′ in 1:NGroups)
            end 
            qig += fixed_sources[i,g]
            qig /= (4π * Σt[g])
            qig = qig < MIN_q ? MIN_q : qig
            q[ig] = qig
        end
    end

    return nothing
end

You need to add a dependency for LinearAlgebra.jl and then import dot from there. We could speed up the fissionable sum as well in a similar manner, but that's not called in this test case.

This makes the compute_q! time comparable to the compute_φ! time, so eliminating the type instability in the latter function would have a more significant impact on the speed of the whole solve.

[orso82]

Nice work @bclyons12! Can't wait to see the deterministic neutronics package in action with these updates!

[bclyons12]

There's also a performance hit in compute_q! in that fixed_sources is accessed in row-major order, so you're skipping around in memory a bunch. It looks difficult to flip the loops around, but you can remove the addition of fixed_sources from those nested loops, then do the loops again in the opposite order and add it in. The function below gave a ~20% speedup on compute_q!, but it's getting to the point where compute_φ! is dominating:

function compute_q!(sol::MoCSolution{T}, prob::MoCProblem, fixed_sources::Matrix{T}) where {T}
    NGroups = ngroups(prob)
    NRegions = nregions(prob)
    @unpack keff, φ, q = sol
    φtmp = Vector{T}(undef, NGroups)

    @inbounds for i in 1:NRegions
        xs = getxs(prob, i)
        @unpack χ, Σt, νΣf, Σs0 = xs
        fissionable = isfissionable(xs)

        for g′ in 1:NGroups
            φtmp[g′] = φ[@region_index(i, g′)]
        end

        for g in 1:NGroups
            ig = @region_index(i, g)
            @views cΣs0 = Σs0[:, g]
            q[ig] = dot(cΣs0, φtmp)
            if fissionable
                q[ig] += sum(1 / keff * χ[g] * νΣf[g′] * φtmp[g′] for g′ in 1:NGroups)
            end
        end
    end

    @inbounds for g in 1:NGroups
        for i in 1:NRegions
            ig = @region_index(i, g)
            q[ig] += fixed_sources[i, g]
            
            xs = getxs(prob, i)
            @unpack χ, Σt, νΣf, Σs0 = xs
            q[ig] /= 4π * Σt[g]

            (q[ig] < MIN_q) && (q[ig] = MIN_q)
        end
    end

    return nothing
end

[kevinm387]

Thank you @bclyons12, this looks like major improvement! I plan to dive back into this later today or tomorrow

[kevinm387]

Mention PR #389

[rvignolo]

Wow, I hope I had the time to improve my NeutronTransport.jl package. It is very nice to see that someone has used it.

Thanks!