Type declarations of Box

[edwinb-ai]

Hello, thanks a lot for the package. It is really fantastic.

I am currently trying to implement a basic Monte Carlo simulation engine here (https://github.com/edwinb-ai/Metropolis.jl), and I am using CellListMap.jl for computing the interactions between particles.

However, I am having some trouble when defining a type that contains a field that is of type Box. The PR in which I am working is here.
Basically, the problem is that I want to define the following type

mutable struct System{B,T,VT,I}
    xpos::VT    # position of particles
    density::T   # density of the system
    temperature::T    # temperature
    box::B    # simulation box, actually want it to be of type `Box`
    rng::Random.AbstractRNG
    npart::I    # total number of particles
end

I then have a constructor that automatically creates the simulation box, which is this one

function System(
    density::T, temp::T, particles::I, cutoff::T; dims=3, random_init=true, lcell=2
) where {T<:Real,I<:Int}
    box_size = cbrt(T(particles) / density)
    # box = create_box(box_size, dims, cutoff; lcell=lcell)
    box = CellListMap.Box(fill(box_size, dims), cutoff; lcell=lcell)
    rng = Xorshifts.Xoroshiro128Plus()
    xpos = initialize_positions(box_size, rng, particles; random_init=random_init)
    syst = System(xpos, density, temp, box, rng, particles)

    return syst
end

I am using Cthulhu.jl to check for type instabilities, and the one that I get is that the type of box is not concrete, and so there are many allocations happening.
I was wondering if you could point me in the right direction of how I should be declaring the types of Box inside other types, such as the one I described above.

[lmiq]

Hello,

Indeed, the "Box" type of CellListMap is annoyingly complex:

Base.@kwdef struct Box{UnitCellType,N,T,TSQ,M}
    unit_cell::UnitCell{UnitCellType,N,T,M}
    lcell::Int
    nc::SVector{N,Int}
    cutoff::T
    cutoff_sq::TSQ
    ranges::SVector{N,UnitRange{Int}}
    cell_size::SVector{N,T}
    unit_cell_max::SVector{N,T}
end

You either need to parameterize all that in your System type, or just use System{B,...} as you are already using. This is what I would recommend really.

The creation of a Box is, however, intrinsically type-unstable, because the UnitCellType is determined at runtime, thus your constructor will be type-unstable as well. The resulting Box is, nevertheless, concrete, so everything from there on should be fast.

The construction of the Box type is certainly irrelevant, in terms of time, in a simulation, thus that type-instability is "benign".

In summary:

julia> using CellListMap

julia> struct Test{B}
           b::B
       end

julia> test = Test(Box([1,1,1],0.1))
Test{Box{OrthorhombicCell, 3, Float64, 9}}(Box{OrthorhombicCell, 3, Float64, 9}(CellListMap.UnitCell{OrthorhombicCell, 3, Float64, 9}([1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]), 1, [12, 12, 12], 0.1, 0.010000000000000002, UnitRange{Int64}[-1:1, -1:1, -1:1], [0.1, 0.1, 0.1], [1.0, 1.0, 1.0]))

julia> typeof(test)
Test{Box{OrthorhombicCell, 3, Float64, 9}}

julia> isconcretetype(typeof(test))
true

The Box constructor is type-unstable, and you cannot get rid of that (if you update the box at every iteration, there will be some allocations associated to that, but very likely in terms of performance that will be irrelevant. It it is not, please let us discuss this further).

julia> @code_warntype Box([1,1,1],0.1)
Variables
  #self#::Type{Box}
  sides::Vector{Int64}
  cutoff::Float64

Body::Box{OrthorhombicCell, _A, Float64, _B} where {_A, _B}
1 ─ %1 = CellListMap.:(var"#Box#12")(CellListMap.Float64, 1, CellListMap.OrthorhombicCell, #self#, sides, cutoff)::Box{OrthorhombicCell, _A, Float64, _B} where {_A, _B}
└──      return %1

julia> @btime Box($([1,1,1]),0.1);
  12.256 μs (101 allocations: 5.62 KiB)

[lmiq]

An additional comment: I noticed that you are trying to pass all parameters to the Box type here:

https://github.com/edwinb-ai/Metropolis.jl/blob/072ba4fcca3e65716b9abfe71a4f28f77bed0020/src/types.jl#L23

That looks fine, though (as you have seen above), the set of type parameters is not something that I guarantee to be stable in the interface. In particular, it will not be in the next release (0.7) because I needed to change them to allow automatic differentiation and unit propagation to happen through the code.

Thus, before you get a broken code yourself there, I really recommend sticking with the B parameter for the complete Box type.

[lmiq]

Just reinforcing: you shouldn't really care that your System constructor is type-stable. The important there is that the resulting System object is concrete, such that the simulation itself using the values contained there is type-stable.

[edwinb-ai]

Oh okay, I see then. Thank you very much for clarifying this. I actually took a lot of inspiration from your FortranCon talk on molecular dynamics, so all of this is greatly appreciated. I still have some other type instabilities, but I'll try to figure those out as I go. Thanks again!

The important there is that the resulting System object is concrete, such that the simulation itself using the values contained there is type-stable.

Thanks, I will make sure to make it concrete.