diff --git a/src/b1pcomponent.jl b/src/b1pcomponent.jl index 2bb48e0..4bd8cbd 100644 --- a/src/b1pcomponent.jl +++ b/src/b1pcomponent.jl @@ -1,20 +1,20 @@ # -# - split off the linear transformation -# - remove the abstract oneparticle basis crap and incorporate it into +# - split off the linear transformation +# - remove the abstract oneparticle basis crap and incorporate it into # the Product1pBasis -> renamed as OnepBasis -# - spec simply specifies in what order the wrapped basis generates the -# basis functions, basically don't allow doing it lazy anymore -# - P_qa(r) is actually P_q(r, a) so the only issue with this is -# if we want different basis length for different species, but even then +# - spec simply specifies in what order the wrapped basis generates the +# basis functions, basically don't allow doing it lazy anymore +# - P_qa(r) is actually P_q(r, a) so the only issue with this is +# if we want different basis length for different species, but even then # we could just fill in zeros # - Linear1pTransform: just a linear mapping # -# MORE NOTES +# MORE NOTES # - consider leaving FVAL in here for now, but move it out later... # - changed getval(fval, X) to evaluate(fval, X) -# - can we get away with only proviging gradients w.r.t. the -# symbols used in this component, or do we need the whole thing? +# - can we get away with only proviging gradients w.r.t. the +# symbols used in this component, or do we need the whole thing? import NamedTupleTools @@ -22,49 +22,49 @@ using NamedTupleTools: namedtuple @doc raw""" -`struct B1pComponent:` Wraps a one-particle basis component. +`struct B1pComponent:` Wraps a one-particle basis component. -`basis` is a structure that can compute a basis, i.e. a vector of scalars, -real or complex. each element of this vector is "specified" by a NamedTuple, -stored in `spec`. E.g., for an `Rn` type basis the spec is would be +`basis` is a structure that can compute a basis, i.e. a vector of scalars, +real or complex. each element of this vector is "specified" by a NamedTuple, +stored in `spec`. E.g., for an `Rn` type basis the spec is would be ```julia spec = [ (n = 1,), (n=2, ), ... ] ``` -while for a Ylm type basis it would be -```julia +while for a Ylm type basis it would be +```julia spec = [ (l=0, m=0), (l=1, m=-1), (l=1, m=0), ... ] ``` -There is no convenience constructor, the `B1pComponent` must be constructed +There is no convenience constructor, the `B1pComponent` must be constructed "by hand", for which several wrapper functions are written. """ -struct B1pComponent{ISYMS, TT, TB, FVAL} - basis::TB - fval::FVAL - spec::Vector{NamedTuple{ISYMS, TT}} - degrees::Vector{Int} - label::String - meta::Dict{String,Any} - # ------------ derived fields - invspec::Dict{NamedTuple{ISYMS, TT}, Int} - # todo - fields for temporary arrays ... +struct B1pComponent{ISYMS,TT,TB,FVAL} + basis::TB + fval::FVAL + spec::Vector{NamedTuple{ISYMS,TT}} + degrees::Vector{Int} + label::String + meta::AbstractDict{String,Any} + # ------------ derived fields + invspec::AbstractDict{NamedTuple{ISYMS,TT},Int} + # todo - fields for temporary arrays ... end -function B1pComponent(basis, fval, spec::AbstractVector{<: NamedTuple}, - degrees::Vector{Int}, - label::AbstractString, - meta = Dict{String, Any}()) - spec1 = collect(spec) - invspec = Dict{eltype(spec1), Int}() - for (i, b) in enumerate(spec) - invspec[b] = i - end - return B1pComponent(basis, fval, spec, degrees, label, meta, invspec) +function B1pComponent(basis, fval, spec::AbstractVector{<:NamedTuple}, + degrees::Vector{Int}, + label::AbstractString, + meta=Dict{String,Any}()) + spec1 = collect(spec) + invspec = Dict{eltype(spec1),Int}() + for (i, b) in enumerate(spec) + invspec[b] = i + end + return B1pComponent(basis, fval, spec, degrees, label, meta, invspec) end -# -------------- management of the basis specification, in particular the -# interaction with Product1pBasis +# -------------- management of the basis specification, in particular the +# interaction with Product1pBasis getlabel(basis::B1pComponent) = basis.label @@ -74,144 +74,142 @@ _idxsyms(basis::B1pComponent{ISYMS}) where {ISYMS} = ISYMS get_spec(basis::B1pComponent) = copy(basis.spec) -get_spec(basis::B1pComponent, i::Integer) = basis.spec[i] +get_spec(basis::B1pComponent, i::Integer) = basis.spec[i] -# TODO - I don't think this is needed anymore, but the sparsification will +# TODO - I don't think this is needed anymore, but the sparsification will # have to be moved in here and something be done about that... # function set_spec!(basis::XScal1pBasis, spec) # basis.spec = identity.(spec) # basis.coeffs = zeros(eltype(basis.coeffs), length(spec), length(basis.P)) -# return basis +# return basis # end -function isadmissible(b::NamedTuple{BSYMS}, basis::B1pComponent) where {BSYMS} - ISYMS = _idxsyms(basis) - # this is an assert since it should ALWAYS be true, if not there is a bug - @assert all(sym in BSYMS for sym in ISYMS) - # project to the ISYMS and check it is in the b1pcomponent - return haskey(basis.invspec, b[ISYMS]) +function isadmissible(b::NamedTuple{BSYMS}, basis::B1pComponent) where {BSYMS} + ISYMS = _idxsyms(basis) + # this is an assert since it should ALWAYS be true, if not there is a bug + @assert all(sym in BSYMS for sym in ISYMS) + # project to the ISYMS and check it is in the b1pcomponent + return haskey(basis.invspec, b[ISYMS]) end -# TODO - LATER +# TODO - LATER # function sparsify!(basis::XScal1pBasis, spec) -# # spec is the part of the basis we keep +# # spec is the part of the basis we keep # inds = Vector{Int}(undef, length(spec)) # for (ib, b) in enumerate(spec) # iold = findall(isequal(b), basis.spec) # @assert length(iold) == 1 -# inds[ib] = iold[1] +# inds[ib] = iold[1] # end # # keep the original order of the basis since we assume it is ordered -# # by some sensible notion of degree. +# # by some sensible notion of degree. # p = sortperm(inds) # basis.spec = spec[p] # basis.coeffs = basis.coeffs[p, :] # return basis # end -symbols(basis::B1pComponent) = [ _idxsyms(basis)... ] +symbols(basis::B1pComponent) = [_idxsyms(basis)...] -function indexrange(basis::B1pComponent) - ISYMS = _idxsyms(basis) - minidx = Dict([sym => minimum(b[sym] for b in basis.spec) for sym in ISYMS]...) - maxidx = Dict([sym => maximum(b[sym] for b in basis.spec) for sym in ISYMS]...) - return Dict([sym => minidx[sym]:maxidx[sym] for sym in ISYMS]...) +function indexrange(basis::B1pComponent) + ISYMS = _idxsyms(basis) + minidx = Dict([sym => minimum(b[sym] for b in basis.spec) for sym in ISYMS]...) + maxidx = Dict([sym => maximum(b[sym] for b in basis.spec) for sym in ISYMS]...) + return Dict([sym => minidx[sym]:maxidx[sym] for sym in ISYMS]...) end -# this is needed to generate the product 1p basis - it returns the -# index of the 1p basis function component used in the 1p basis function -# specified by b -function get_index(basis::B1pComponent, b::NamedTuple) - ISYMS = _idxsyms(basis) - b1 = b[ISYMS] - if !haskey(basis.invspec, b1) - error("B1pComponent ($(basis.label) : can't find $(b1) in spec") - end - return basis.invspec[b1] +# this is needed to generate the product 1p basis - it returns the +# index of the 1p basis function component used in the 1p basis function +# specified by b +function get_index(basis::B1pComponent, b::NamedTuple) + ISYMS = _idxsyms(basis) + b1 = b[ISYMS] + if !haskey(basis.invspec, b1) + error("B1pComponent ($(basis.label) : can't find $(b1) in spec") + end + return basis.invspec[b1] end # -------------- degree calculations -# this should be revisited, doesn't feel clean yet +# this should be revisited, doesn't feel clean yet function degree(b::NamedTuple, basis::B1pComponent) - ISYMS = _idxsyms(basis) - idx = basis.invspec[ b[ISYMS] ] - return basis.degrees[idx] + ISYMS = _idxsyms(basis) + idx = basis.invspec[b[ISYMS]] + return basis.degrees[idx] end -function degree(b::NamedTuple, basis::B1pComponent, weight::Dict) - ISYMS = _idxsyms(basis) - return weight[ISYMS[1]] * degree(b, basis) +function degree(b::NamedTuple, basis::B1pComponent, weight::AbstractDict) + ISYMS = _idxsyms(basis) + return weight[ISYMS[1]] * degree(b, basis) end -# --------------- IO operations +# --------------- IO operations -==(P1::B1pComponent, P2::B1pComponent) = ( - (P1.basis == P2.basis) && (P1.spec == P2.spec) && - (P1.label == P2.label) && (P1.fval == P2.fval) && - (P1.degrees == P2.degrees) ) +==(P1::B1pComponent, P2::B1pComponent) = ( + (P1.basis == P2.basis) && (P1.spec == P2.spec) && + (P1.label == P2.label) && (P1.fval == P2.fval) && + (P1.degrees == P2.degrees)) function write_dict(basis::B1pComponent) - ISYMS = _idxsyms(basis) - return Dict("__id__" => "ACEfrictionCore_B1pComponent", - "syms" => [ string.(ISYMS) ...], - "basis" => write_dict(basis.basis), - "fval" => write_dict(basis.fval), - "spec" => convert.(Dict, basis.spec), - "degrees" => basis.degrees, - "label" => basis.label ) + ISYMS = _idxsyms(basis) + return Dict("__id__" => "ACEfrictionCore_B1pComponent", + "syms" => [string.(ISYMS)...], + "basis" => write_dict(basis.basis), + "fval" => write_dict(basis.fval), + "spec" => convert.(Dict, basis.spec), + "degrees" => basis.degrees, + "label" => basis.label) end -function read_dict(::Val{:ACEfrictionCore_B1pComponent}, D::Dict) - basis = read_dict(D["basis"]) - ISYMS = tuple(Symbol.(D["syms"])...) - spec = NamedTuple{ISYMS}.(namedtuple.(D["spec"])) - fval = read_dict(D["fval"]) - return B1pComponent(basis, fval, spec, Int.(D["degrees"]), D["label"]) +function read_dict(::Val{:ACEfrictionCore_B1pComponent}, D::AbstractDict) + basis = read_dict(D["basis"]) + ISYMS = tuple(Symbol.(D["syms"])...) + spec = NamedTuple{ISYMS}.(namedtuple.(D["spec"])) + fval = read_dict(D["fval"]) + return B1pComponent(basis, fval, spec, Int.(D["degrees"]), D["label"]) end function show(io::IO, basis::B1pComponent) - vsyms = Transforms.get_symbols(basis.fval) - strvsyms = filter(!isequal(':'), "$vsyms") - if length(vsyms) == 1 - # this doesn't work with unicode!?!?!? - # strvsyms = strvsyms[1:end-2] * ")" - strvsyms = prod(split(strvsyms, ',')) - end - print(io, "B1pComponent( $(basis.label)$(strvsyms)") - for (key, rg) in indexrange(basis) - print(io, ", $key = $(rg)") - end - print(io, ")") + vsyms = Transforms.get_symbols(basis.fval) + strvsyms = filter(!isequal(':'), "$vsyms") + if length(vsyms) == 1 + # this doesn't work with unicode!?!?!? + # strvsyms = strvsyms[1:end-2] * ")" + strvsyms = prod(split(strvsyms, ',')) + end + print(io, "B1pComponent( $(basis.label)$(strvsyms)") + for (key, rg) in indexrange(basis) + print(io, ", $key = $(rg)") + end + print(io, ")") end # ------------------------ Evaluation code # this is basically an interface for the inner basis -evaluate(basis::B1pComponent, X::AbstractState) = - evaluate(basis.basis, evaluate(basis.fval, X)) +evaluate(basis::B1pComponent, X::AbstractState) = + evaluate(basis.basis, evaluate(basis.fval, X)) # this one we probably only need for training so ... -# - we can relax the efficiency a bit -# - we actually never need the evaluate_dd but only the -# associated backpropagation operation which may be a bit simpler -# to implement. +# - we can relax the efficiency a bit +# - we actually never need the evaluate_dd but only the +# associated backpropagation operation which may be a bit simpler +# to implement. # so we will likely remove this entirely. # evaluate_dd(basis::B1pComponent, X::AbstractState) -# -------------- AD codes - - +# -------------- AD codes diff --git a/src/basisselectors.jl b/src/basisselectors.jl index 3e13529..adfd4bc 100644 --- a/src/basisselectors.jl +++ b/src/basisselectors.jl @@ -1,71 +1,71 @@ -# Interface +# Interface -const Onepb = NamedTuple -const Prodb = AbstractVector{<: NamedTuple} +const Onepb = NamedTuple +const Prodb = AbstractVector{<:NamedTuple} """ `AbstractBasisSelector` : object specifying how a finite basis is selected from -the infinite space of symmetric polynomials. This type is pretty superfluous -for now since all basis selector and basis selection algorithms we admit so -far require Downsets. +the infinite space of symmetric polynomials. This type is pretty superfluous +for now since all basis selector and basis selection algorithms we admit so +far require Downsets. """ abstract type AbstractBasisSelector end """ -`DownsetBasisSelector` : must implemented a non-negative valued `level` function -that is monotonically increasing with respect to the lexiographic ordering of +`DownsetBasisSelector` : must implemented a non-negative valued `level` function +that is monotonically increasing with respect to the lexiographic ordering of the basis functions. """ abstract type DownsetBasisSelector <: AbstractBasisSelector end """ -return maximum correlation order of the basis selector. Note that the filter +return maximum correlation order of the basis selector. Note that the filter function could reduce that by a bit. """ -function maxorder end +function maxorder end """ * `level(b::Onepb, Bsel::DownsetBasisSelector, basis::OneParticleBasis)` * `level(bb::Prodb, Bsel::DownsetBasisSelector, basis::OneParticleBasis)` -The first version specifies the level function for +The first version specifies the level function for """ -function level end +function level end """ `maxlevel(Bsel, basis1p)` """ -function maxlevel end +function maxlevel end """ `level1(b::Onepb, Bsel::DownsetBasisSelector, basis::OneParticleBasis)` -A specialized version of `level` to construct the 1-particle basis (cf. -`init1pspec!`). Fallback is to just call the `level` function, but this gives -some additional flexibility to ensure that the conditions of the `level` -framework are all justified, see docs for more detail. +A specialized version of `level` to construct the 1-particle basis (cf. +`init1pspec!`). Fallback is to just call the `level` function, but this gives +some additional flexibility to ensure that the conditions of the `level` +framework are all justified, see docs for more detail. """ -level1(b::Onepb, Bsel::DownsetBasisSelector, basis::OneParticleBasis) = - level(b, Bsel, basis) +level1(b::Onepb, Bsel::DownsetBasisSelector, basis::OneParticleBasis) = + level(b, Bsel, basis) -maxlevel1(Bsel::DownsetBasisSelector, basis::OneParticleBasis) = - maxlevel(Bsel, basis) +maxlevel1(Bsel::DownsetBasisSelector, basis::OneParticleBasis) = + maxlevel(Bsel, basis) """ `filter(b_or_bb, Bsel::AbstractBasisSelector, basis::OneParticleBasis)` -After a down-set basis has been constructed, it can still be filtered, which -allows us to construct basis sets that aren't downsets but not too far away from -downsets. The main application is to enfore the constraints on m and l channels -arising from the symmetries. +After a down-set basis has been constructed, it can still be filtered, which +allows us to construct basis sets that aren't downsets but not too far away from +downsets. The main application is to enfore the constraints on m and l channels +arising from the symmetries. -Fallback implementation always returns `true`. +Fallback implementation always returns `true`. """ filter(b_or_bb, Bsel::AbstractBasisSelector, basis::OneParticleBasis) = true @@ -73,20 +73,20 @@ filter(b_or_bb, Bsel::AbstractBasisSelector, basis::OneParticleBasis) = true """ No constraints on the basis - this selects that largest possible basis subject to additional constraints baked into the one-particle basis. -In practise this should be equivalent to a naive max-norm basis selection, -and likely never used in practise. +In practise this should be equivalent to a naive max-norm basis selection, +and likely never used in practise. """ struct MaxBasis <: DownsetBasisSelector - maxorder::Int + maxorder::Int end -maxorder(Bsel::MaxBasis) = Bsel.maxorder +maxorder(Bsel::MaxBasis) = Bsel.maxorder -level(b::Onepb, Bsel::MaxBasis, basis::OneParticleBasis) = - degree(b, basis) +level(b::Onepb, Bsel::MaxBasis, basis::OneParticleBasis) = + degree(b, basis) -level(bb::Prodb, Bsel::MaxBasis, basis::OneParticleBasis) = - length(bb) == 0 ? 0 : sum(b -> level(b, Bsel, basis), bb) +level(bb::Prodb, Bsel::MaxBasis, basis::OneParticleBasis) = + length(bb) == 0 ? 0 : sum(b -> level(b, Bsel, basis), bb) maxlevel(Bsel::MaxBasis, basis::OneParticleBasis) = Inf @@ -97,44 +97,44 @@ maxlevel(Bsel::MaxBasis, basis::OneParticleBasis) = Inf The most basic form of a sparse basis selection, using the total degree. Only the maximum correlation order and maximum degree may be specified. -This should primarily be used for testing. Construct it via +This should primarily be used for testing. Construct it via ```julia SimpleSparseBasis(maxorder, maxdegree) ``` -Note that `maxdegree` is a basic total degree. +Note that `maxdegree` is a basic total degree. """ struct SimpleSparseBasis <: DownsetBasisSelector - maxorder::Int - maxlevel::Float64 + maxorder::Int + maxlevel::Float64 end level(b::Onepb, Bsel::SimpleSparseBasis, basis::OneParticleBasis) = - degree(b, basis) + degree(b, basis) level(bb::Prodb, Bsel::SimpleSparseBasis, basis::OneParticleBasis) = - length(bb) == 0 ? 0 : sum( degree(b, basis) for b in bb ) + length(bb) == 0 ? 0 : sum(degree(b, basis) for b in bb) -maxlevel(bb, Bsel::SimpleSparseBasis, args...) = - Bsel.maxlevel +maxlevel(bb, Bsel::SimpleSparseBasis, args...) = + Bsel.maxlevel -maxlevel(Bsel::SimpleSparseBasis, args...) = - Bsel.maxlevel +maxlevel(Bsel::SimpleSparseBasis, args...) = + Bsel.maxlevel -maxorder(Bsel::SimpleSparseBasis, args...) = - Bsel.maxorder +maxorder(Bsel::SimpleSparseBasis, args...) = + Bsel.maxorder """ `AbstractSparseBasis`: Super-type for sparse basis selection as sub-levelsets of the levelset function `level` and corresponding (possibly order-dependent) -levels provided in the dictionary `maxlevels::Dict{Any, Float64}`. In the default +levels provided in the dictionary `maxlevels::AbstractDict{Any, Float64}`. In the default implementation the levelset function and the degree function are identical. -Basis functions are selected in two steps. First, "admissible" basis -specifications are generated as a sub-levelset of the leveset function using the +Basis functions are selected in two steps. First, "admissible" basis +specifications are generated as a sub-levelset of the leveset function using the implementation of the function `gensparse`. After that, basis functions that do -not satisfy the conditons implemented in the function `filter` are removed from +not satisfy the conditons implemented in the function `filter` are removed from the basis set. """ abstract type AbstractSparseBasis <: DownsetBasisSelector end @@ -142,107 +142,107 @@ abstract type AbstractSparseBasis <: DownsetBasisSelector end maxorder(Bsel::AbstractSparseBasis) = Bsel.maxorder level(b::Onepb, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = - degree(b, basis, Bsel.weight) + degree(b, basis, Bsel.weight) -level(bb::Prodb, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = ( - length(bb) == 0 ? 0.0 - : norm( level.(bb, Ref(Bsel), Ref(basis)), Bsel.p ) ) +level(bb::Prodb, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = ( + length(bb) == 0 ? 0.0 + : norm(level.(bb, Ref(Bsel), Ref(basis)), Bsel.p)) -maxlevel(bb::Prodb, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = - maxlevel(length(bb), Bsel, basis) +maxlevel(bb::Prodb, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = + maxlevel(length(bb), Bsel, basis) -maxlevel(ord::Integer, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = - ( haskey(Bsel.maxlevels, ord) ? Bsel.maxlevels[ord] - : Bsel.maxlevels["default"] ) +maxlevel(ord::Integer, Bsel::AbstractSparseBasis, basis::OneParticleBasis) = + (haskey(Bsel.maxlevels, ord) ? Bsel.maxlevels[ord] + : Bsel.maxlevels["default"]) -maxlevel(Bsel::AbstractSparseBasis, basis::OneParticleBasis) = - maximum( maxlevel(ord, Bsel, basis) for ord = 1:maxorder(Bsel) ) +maxlevel(Bsel::AbstractSparseBasis, basis::OneParticleBasis) = + maximum(maxlevel(ord, Bsel, basis) for ord = 1:maxorder(Bsel)) """ `SparseBasis`: basic implementation of an `AbstractSparseBasis`. -- `maxorder` : maximum correlation order -- `p` : degree (level) is computed via a weighted p-norm of the 1p basis +- `maxorder` : maximum correlation order +- `p` : degree (level) is computed via a weighted p-norm of the 1p basis function degrees -- `maxlevels` : for each order `ord`, `maxlevels[ord]` is the maximum level - for basis functions of that order. if `maxlevels[ord]` is not +- `maxlevels` : for each order `ord`, `maxlevels[ord]` is the maximum level + for basis functions of that order. if `maxlevels[ord]` is not specified, then `maxlevels["default"]` is the fallback. -- `weight` : specify weights for the different channels of the 1p basis - functions, e.g., for an `Rn * Ylm` basis, one might specify - `weight[:n] = 1` and `weight[:l] = 2` to have more radial and - fewer angular basis functions. +- `weight` : specify weights for the different channels of the 1p basis + functions, e.g., for an `Rn * Ylm` basis, one might specify + `weight[:n] = 1` and `weight[:l] = 2` to have more radial and + fewer angular basis functions. """ struct SparseBasis <: AbstractSparseBasis - maxorder::Int - weight::Dict{Symbol, Float64} - maxlevels::Dict{Any, Float64} - p::Float64 + maxorder::Int + weight::AbstractDict{Symbol,Float64} + maxlevels::AbstractDict{Any,Float64} + p::Float64 end -function SparseBasis(; maxorder::Integer = nothing, - p = 1, - weight = Dict(:l => 1.0, :n => 1.0), - default_maxdeg = nothing, - maxlevels = nothing ) - if (default_maxdeg != nothing) && (maxlevels == nothing ) - return SparseBasis(maxorder, weight, - Dict{Any, Float64}("default" => default_maxdeg), - p) - elseif (default_maxdeg == nothing) && (maxlevels != nothing) - SparseBasis(maxorder, weight, maxlevels, p) - else - error("""Either both or neither optional arguments `maxlevels` and - `default_maxdeg` were provided. To avoid ambiguity ensure that - exactly one of these arguments is provided.""") - end +function SparseBasis(; maxorder::Integer=nothing, + p=1, + weight=Dict(:l => 1.0, :n => 1.0), + default_maxdeg=nothing, + maxlevels=nothing) + if (default_maxdeg != nothing) && (maxlevels == nothing) + return SparseBasis(maxorder, weight, + Dict{Any,Float64}("default" => default_maxdeg), + p) + elseif (default_maxdeg == nothing) && (maxlevels != nothing) + SparseBasis(maxorder, weight, maxlevels, p) + else + error("""Either both or neither optional arguments `maxlevels` and + `default_maxdeg` were provided. To avoid ambiguity ensure that + exactly one of these arguments is provided.""") + end end """ -`CategorySparseBasis`: extension of `SparseBasis` that implements a +`CategorySparseBasis`: extension of `SparseBasis` that implements a constraint on the correlation orders for within-category correlations, i.e., for -each category `cat` contained in OneParticleBasis associated with the symbol +each category `cat` contained in OneParticleBasis associated with the symbol `isym`, it is required that the within-category correlation order `corr[cat]` satisfies -```julia +```julia minorder_dict[cat] <= corr[cat] <= maxorder_dict[cat]. ``` """ struct CategorySparseBasis <: AbstractSparseBasis - maxorder::Int - weight::Dict{Symbol, Float64} - maxlevels::Dict{Any, Float64} - p::Float64 - isym::Symbol - minorder_dict::Dict{Any, Int} - maxorder_dict::Dict{Any, Int} - weight_cat::Dict{Any, Float64} + maxorder::Int + weight::AbstractDict{Symbol,Float64} + maxlevels::AbstractDict{Any,Float64} + p::Float64 + isym::Symbol + minorder_dict::AbstractDict{Any,Int} + maxorder_dict::AbstractDict{Any,Int} + weight_cat::AbstractDict{Any,Float64} end -function CategorySparseBasis(isym::Symbol, categories; - maxorder::Integer = nothing, - p = 1, - weight = Dict{Symbol, Float64}(), - default_maxdeg = nothing, - maxlevels::Dict{Any, Float64} = nothing, - minorder_dict = Dict{Any, Float64}(), - maxorder_dict = Dict{Any, Float64}(), - weight_cat = Dict(c => 1.0 for c in categories), - ) - if (default_maxdeg != nothing) && (maxlevels == nothing ) - CategorySparseBasis(maxorder, weight, - Dict{Any, Float64}("default" => default_maxdeg), - p, isym, minorder_dict, maxorder_dict, weight_cat) - elseif (default_maxdeg == nothing) && (maxlevels != nothing) - CategorySparseBasis(maxorder, weight, maxlevels, p, isym, minorder_dict, - maxorder_dict, weight_cat) - else - @error """Either both or neither optional arguments `maxlevels` and - `default_maxdeg` were provided. To avoid ambiguity ensure that - exactly one of these arguments is provided.""" - end +function CategorySparseBasis(isym::Symbol, categories; + maxorder::Integer=nothing, + p=1, + weight=AbstractDict{Symbol,Float64}(), + default_maxdeg=nothing, + maxlevels::AbstractDict{Any,Float64}=nothing, + minorder_dict=AbstractDict{Any,Float64}(), + maxorder_dict=AbstractDict{Any,Float64}(), + weight_cat=Dict(c => 1.0 for c in categories), +) + if (default_maxdeg != nothing) && (maxlevels == nothing) + CategorySparseBasis(maxorder, weight, + Dict{Any,Float64}("default" => default_maxdeg), + p, isym, minorder_dict, maxorder_dict, weight_cat) + elseif (default_maxdeg == nothing) && (maxlevels != nothing) + CategorySparseBasis(maxorder, weight, maxlevels, p, isym, minorder_dict, + maxorder_dict, weight_cat) + else + @error """Either both or neither optional arguments `maxlevels` and + `default_maxdeg` were provided. To avoid ambiguity ensure that + exactly one of these arguments is provided.""" + end end maxorder(Bsel::CategorySparseBasis, category) = Bsel.maxorder_dict[category] @@ -251,69 +251,69 @@ minorder(Bsel::CategorySparseBasis, category) = Bsel.minorder_dict[category] filter(b::ACEfrictionCore.Onepb, Bsel::CategorySparseBasis, basis::OneParticleBasis) = true -function filter(bb, Bsel::CategorySparseBasis, basis::OneParticleBasis) - # auxiliary function to count the number of 1pbasis functions in bb - # for which b.isym == s. - num_b_is_(s) = sum([(getproperty(b, Bsel.isym) == s) for b in bb]) +function filter(bb, Bsel::CategorySparseBasis, basis::OneParticleBasis) + # auxiliary function to count the number of 1pbasis functions in bb + # for which b.isym == s. + num_b_is_(s) = sum([(getproperty(b, Bsel.isym) == s) for b in bb]) - # Within category min correlation order constaint: - cond_ord_cats_min = all( num_b_is_(s) >= minorder(Bsel, s) - for s in keys(Bsel.minorder_dict) ) - # Within category max correlation order constaint: - cond_ord_cats_max = all( num_b_is_(s) <= maxorder(Bsel, s) - for s in keys(Bsel.maxorder_dict) ) + # Within category min correlation order constaint: + cond_ord_cats_min = all(num_b_is_(s) >= minorder(Bsel, s) + for s in keys(Bsel.minorder_dict)) + # Within category max correlation order constaint: + cond_ord_cats_max = all(num_b_is_(s) <= maxorder(Bsel, s) + for s in keys(Bsel.maxorder_dict)) - return cond_ord_cats_min && cond_ord_cats_max + return cond_ord_cats_min && cond_ord_cats_max end -# maxorder and maxlevel are inherited from the abstract interface +# maxorder and maxlevel are inherited from the abstract interface level(b::Onepb, Bsel::CategorySparseBasis, basis::OneParticleBasis) = - cat_weighted_degree(b, Bsel, basis) + cat_weighted_degree(b, Bsel, basis) -level(bb::Prodb, Bsel::CategorySparseBasis, basis::OneParticleBasis) = - cat_weighted_degree(bb, Bsel, basis) +level(bb::Prodb, Bsel::CategorySparseBasis, basis::OneParticleBasis) = + cat_weighted_degree(bb, Bsel, basis) # Category-weighted degree function cat_weighted_degree(b::Onepb, Bsel::CategorySparseBasis, basis::OneParticleBasis) = - degree(b, basis, Bsel.weight) * Bsel.weight_cat[getproperty(b, Bsel.isym)] + degree(b, basis, Bsel.weight) * Bsel.weight_cat[getproperty(b, Bsel.isym)] cat_weighted_degree(bb::Prodb, Bsel::CategorySparseBasis, basis::OneParticleBasis) = ( - length(bb) == 0 ? 0.0 - : norm(cat_weighted_degree.(bb, Ref(Bsel), Ref(basis)), Bsel.p) - ) + length(bb) == 0 ? 0.0 + : norm(cat_weighted_degree.(bb, Ref(Bsel), Ref(basis)), Bsel.p) +) -# --------------------------- -# Some useful filters +# --------------------------- +# Some useful filters -struct NoConstant +struct NoConstant end (::NoConstant)(bb) = (length(bb) > 0) """ -`EvenL`: selects all basis functions where the sum `L = sum_i l_i` of the degrees `l_i` of the spherical harmonics is even. +`EvenL`: selects all basis functions where the sum `L = sum_i l_i` of the degrees `l_i` of the spherical harmonics is even. """ struct EvenL - isym::Symbol - categories + isym::Symbol + categories end - -function (f::ACEfrictionCore.EvenL)(bb) - if isempty(bb) - return true - else - suml(s) = sum( [getl(O3(), b) for b in bb if getproperty(b, f.isym) == s]) - return all(iseven(suml(s)) for s in f.categories) - end + +function (f::ACEfrictionCore.EvenL)(bb) + if isempty(bb) + return true + else + suml(s) = sum([getl(O3(), b) for b in bb if getproperty(b, f.isym) == s]) + return all(iseven(suml(s)) for s in f.categories) + end end #= """ -`DownsetIntersection`: Basis selector whose set of admissible specifications is the intersection +`DownsetIntersection`: Basis selector whose set of admissible specifications is the intersection of the sets of admissible specifications of the sparse basis selectors contained in the lists `DBsels` and `ABsels`. """ struct DownsetIntersection <: DownsetBasisSelector @@ -322,7 +322,7 @@ struct DownsetIntersection <: DownsetBasisSelector maxorder::Int end -maxlevel(Bsel::DownsetIntersection, basis::OneParticleBasis) = 1.0 +maxlevel(Bsel::DownsetIntersection, basis::OneParticleBasis) = 1.0 maxorder(Bsel::DownsetIntersection) = Bsel.maxorder @@ -334,7 +334,7 @@ function Base.intersect(Bsel1::DownsetIntersection,Bsel2::AbstractBasisSelector) return DownsetIntersection(Bsel1.DBsels, vcat(Bsel1.ABsels,[Bsel2]), minimum([maxorder(Bsel1),maxorder(Bsel2)])) end -function Base.intersect(Bsel1::DownsetBasisSelector, Bsel2::AbstractBasisSelector) +function Base.intersect(Bsel1::DownsetBasisSelector, Bsel2::AbstractBasisSelector) Bsel = DownsetIntersection(Vector{DownsetBasisSelector}([]), Vector{AbstractBasisSelector}([]), Inf) return intersect(intersect(Bsel, Bsel1), Bsel2) end @@ -350,12 +350,12 @@ function Base.intersect(Bsel1::DownsetIntersection, Bsel2::DownsetIntersection) return Bsel end -function level(bb, Bsel::DownsetIntersection, basis::OneParticleBasis) +function level(bb, Bsel::DownsetIntersection, basis::OneParticleBasis) return maximum([ level(bb, bsel, basis)/maxlevel(bsel,basis) for bsel in Bsel.DBsels]) #return maximum([ level(bb, bsel, basis)/maxlevel(length(bb),bsel,basis) for bsel in Bsel.DBsels]) end -function filter(bb, Bsel::DownsetIntersection, basis::OneParticleBasis) +function filter(bb, Bsel::DownsetIntersection, basis::OneParticleBasis) return all([filter(bb, bsel, basis) for bsel in Bsel.DBsels]) && all([filter(bb, bsel, basis) for bsel in Bsel.ABsels]) end =# diff --git a/src/bonds/bondcutoffs.jl b/src/bonds/bondcutoffs.jl index 20f51df..8ba1427 100644 --- a/src/bonds/bondcutoffs.jl +++ b/src/bonds/bondcutoffs.jl @@ -127,13 +127,13 @@ include("ellipsoid_trans.jl") # contains rrules for ellipsoid2sphere and other a function ACEfrictionCore.write_dict(cutoff::EllipsoidCutoff{T}) where {T} Dict("__id__" => "ACEbonds_EllipsoidCutoff", - "rcutbond" => cutoff.rcutbond, - "rcutenv" => cutoff.rcutenv, - "zcutenv" => cutoff.zcutenv, - "T" => T) -end + "rcutbond" => cutoff.rcutbond, + "rcutenv" => cutoff.rcutenv, + "zcutenv" => cutoff.zcutenv, + "T" => T) +end -function ACEfrictionCore.read_dict(::Val{:ACEbonds_EllipsoidCutoff}, D::Dict) +function ACEfrictionCore.read_dict(::Val{:ACEbonds_EllipsoidCutoff}, D::AbstractDict) rcutbond = D["rcutbond"] rcutenv = D["rcutenv"] zcutenv = D["zcutenv"] diff --git a/src/bonds/bondpot.jl b/src/bonds/bondpot.jl index de49f15..21ea7be 100644 --- a/src/bonds/bondpot.jl +++ b/src/bonds/bondpot.jl @@ -4,67 +4,67 @@ import ACEfrictionCore.ACEbonds.BondCutoffs: env_transform, env_filter, Abstract import ACEfrictionCore: params, nparams, set_params! export params, nparams, set_params! -# TODO: extend implementation to allow for LinearModels with multiple featuers. +# TODO: extend implementation to allow for LinearModels with multiple featuers. struct ACEBondPotential{TM} <: AbstractCalculator - models::Dict{Tuple{AtomicNumber, AtomicNumber}, TM} - cutoff::AbstractBondCutoff{Float64} + models::AbstractDict{Tuple{AtomicNumber,AtomicNumber},TM} + cutoff::AbstractBondCutoff{Float64} end struct ACEBondPotentialBasis{TM} <: JuLIP.MLIPs.IPBasis - models::Dict{Tuple{AtomicNumber, AtomicNumber}, TM} # model = basis - inds::Dict{Tuple{AtomicNumber, AtomicNumber}, UnitRange{Int}} - cutoff::AbstractBondCutoff{Float64} + models::AbstractDict{Tuple{AtomicNumber,AtomicNumber},TM} # model = basis + inds::AbstractDict{Tuple{AtomicNumber,AtomicNumber},UnitRange{Int}} + cutoff::AbstractBondCutoff{Float64} end function basis(V::ACEBondPotential) - models = Dict( [zz => model.basis for (zz, model) in V.models]... ) - inds = _get_basisinds(V) - return ACEBondPotentialBasis(models, inds, V.cutoff) + models = Dict([zz => model.basis for (zz, model) in V.models]...) + inds = _get_basisinds(V) + return ACEBondPotentialBasis(models, inds, V.cutoff) end -ACEBondCalc = Union{ACEBondPotential, ACEBondPotentialBasis} +ACEBondCalc = Union{ACEBondPotential,ACEBondPotentialBasis} -function params(calc::ACEBondPotential) - θ = zeros(nparams(calc)) - inds = _get_basisinds(calc) - for zz in keys(inds) - m = _get_model(calc, zz[1],zz[2]) - θ[inds[zz]] = params(m) - end - return θ +function params(calc::ACEBondPotential) + θ = zeros(nparams(calc)) + inds = _get_basisinds(calc) + for zz in keys(inds) + m = _get_model(calc, zz[1], zz[2]) + θ[inds[zz]] = params(m) + end + return θ end function set_params!(calc::ACEBondPotential, θ) - inds = _get_basisinds(calc) - for zz in keys(calc.models) - ACEfrictionCore.set_params!(calc, zz, θ[inds[zz]]) - end + inds = _get_basisinds(calc) + for zz in keys(calc.models) + ACEfrictionCore.set_params!(calc, zz, θ[inds[zz]]) + end end function set_params!(calc::ACEBondPotential, zz::Tuple{AtomicNumber,AtomicNumber}, θ) - set_params!(_get_model(calc, zz[1], zz[2]),θ) + set_params!(_get_model(calc, zz[1], zz[2]), θ) end nparams(V::ACEBondPotential) = sum(length(inds) for (_, inds) in _get_basisinds(V)) -Base.length(basis::ACEBondPotentialBasis) = - sum(length(inds) for (_, inds) in basis.inds) +Base.length(basis::ACEBondPotentialBasis) = + sum(length(inds) for (_, inds) in basis.inds) function _get_basisinds(V::ACEBondPotential) - inds = Dict{Tuple{AtomicNumber, AtomicNumber}, UnitRange{Int}}() - zz = sort(collect(keys(V.models))) - i0 = 0 - for z in zz - mo = V.models[z] - len = length(mo.basis) - inds[z] = (i0+1):(i0+len) # to generalize for general models - i0 += len - end - return inds + inds = Dict{Tuple{AtomicNumber,AtomicNumber},UnitRange{Int}}() + zz = sort(collect(keys(V.models))) + i0 = 0 + for z in zz + mo = V.models[z] + len = length(mo.basis) + inds[z] = (i0+1):(i0+len) # to generalize for general models + i0 += len + end + return inds end _get_basisinds(V::ACEBondPotentialBasis) = V.inds @@ -72,30 +72,30 @@ _get_basisinds(V::ACEBondPotentialBasis) = V.inds # -------------------------------------------------------- import JuLIP: energy -#, forces, virial +#, forces, virial import ACEfrictionCore: evaluate #, evaluate_d, grad_config -# overload the initiation of the bonds iterator to correctly extract the -# right cutoffs. -bonds(at::Atoms, calc::ACEBondCalc, args...) = bonds(at, calc.cutoff, args...) +# overload the initiation of the bonds iterator to correctly extract the +# right cutoffs. +bonds(at::Atoms, calc::ACEBondCalc, args...) = bonds(at, calc.cutoff, args...) -_get_model(calc::ACEBondCalc, zi, zj) = - calc.models[(min(zi, zj), max(zi,zj))] +_get_model(calc::ACEBondCalc, zi, zj) = + calc.models[(min(zi, zj), max(zi, zj))] function energy(calc::ACEBondPotential, at::Atoms) - E = 0.0 - for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc) - # find the right ace model - ace = _get_model(calc, at.Z[i], at.Z[j]) - # transform the euclidean to cylindrical coordinates - env = env_transform(rrij, at.Z[i], at.Z[j], Rs, Zs, calc.cutoff) - # evaluate - Eij = evaluate(ace, env) - E += Eij.val - end - return E + E = 0.0 + for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc) + # find the right ace model + ace = _get_model(calc, at.Z[i], at.Z[j]) + # transform the euclidean to cylindrical coordinates + env = env_transform(rrij, at.Z[i], at.Z[j], Rs, Zs, calc.cutoff) + # evaluate + Eij = evaluate(ace, env) + E += Eij.val + end + return E end @@ -103,24 +103,24 @@ end # F = zeros(SVector{3, Float64}, length(at)) # for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc) # Zi, Zj = at.Z[i], at.Z[j] -# # find the right ace model +# # find the right ace model # ace = _get_model(calc, Zi, Zj) # # transform the euclidean to cylindrical coordinates # env = env_transform(rrij, Zi, Zj, Rs, Zs, calc.cutoff) -# # evaluate +# # evaluate # dV_cyl = grad_config(ace, env) -# # transform back? +# # transform back? # dV_drrij, dV_dRs = rrule_env_transform(rrij::SVector, Zi, Zj, Rs, Zs, dV_cyl, calc.cutoff) -# # assemble the forces -# F[i] += dV_drrij -# F[j] -= dV_drrij +# # assemble the forces +# F[i] += dV_drrij +# F[j] -= dV_drrij # for (k, dv) in zip(Js, dV_dRs) # F[k] -= dv -# F[i] += 0.5 * dv -# F[j] += 0.5 * dv +# F[i] += 0.5 * dv +# F[j] += 0.5 * dv # end # end -# return F +# return F # end # site_virial(dV::AbstractVector{JVec{T1}}, R::AbstractVector{JVec{T2}} @@ -129,20 +129,20 @@ end # : zero(JMat{fltype_intersect(T1, T2)}) # ) -# function virial(calc::ACEBondPotential, at::Atoms{T}) where {T} +# function virial(calc::ACEBondPotential, at::Atoms{T}) where {T} # V = zero(SMatrix{3, 3, T}) # for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc) # Zi, Zj = at.Z[i], at.Z[j] -# # find the right ace model for this bond +# # find the right ace model for this bond # ace = _get_model(calc, Zi, Zj) # # transform the euclidean to cylindrical coordinates # env = env_transform(rrij, Zi, Zj, Rs, Zs, calc.cutoff) -# # evaluate +# # evaluate # dV_cyl = grad_config(ace, env) -# # transform back? +# # transform back? # dV_drrij, dV_dRs = rrule_env_transform(rrij::SVector, Zi, Zj, Rs, Zs, dV_cyl, calc.cutoff) -# # assemble the virial -# # dV_dRs contain derivative relative to midpoint +# # assemble the virial +# # dV_dRs contain derivative relative to midpoint # # dV_drrij contain derivative w.r.t. rrij # V -= dV_drrij * rrij' # for q = 1:length(dV_dRs) @@ -155,18 +155,18 @@ end function energy(basis::ACEBondPotentialBasis, at::Atoms) - E = zeros(Float64, length(basis)) - Et = zeros(Float64, length(basis)) - for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis) - # find the right ace model - ace = _get_model(basis, at.Z[i], at.Z[j]) - # transform the euclidean to cylindrical coordinates - env = env_transform(rrij, at.Z[i], at.Z[j], Rs, Zs, basis.cutoff) - # evaluate - ACEfrictionCore.evaluate!(Et, ace, ACEfrictionCore.ACEConfig(env)) - E += Et - end - return E + E = zeros(Float64, length(basis)) + Et = zeros(Float64, length(basis)) + for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis) + # find the right ace model + ace = _get_model(basis, at.Z[i], at.Z[j]) + # transform the euclidean to cylindrical coordinates + env = env_transform(rrij, at.Z[i], at.Z[j], Rs, Zs, basis.cutoff) + # evaluate + ACEfrictionCore.evaluate!(Et, ace, ACEfrictionCore.ACEConfig(env)) + E += Et + end + return E end @@ -174,39 +174,39 @@ end # F = zeros(SVector{3, Float64}, length(basis), length(at)) # for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis) # Zi, Zj = at.Z[i], at.Z[j] -# # find the right ace model +# # find the right ace model # ace = _get_model(basis, Zi, Zj) # # transform the euclidean to cylindrical coordinates # env = env_transform(rrij, Zi, Zj, Rs, Zs, basis.cutoff) -# # evaluate +# # evaluate # dB_cyl = evaluate_d(ace, env) -# # transform back? +# # transform back? # dB_drrij, dB_dRs = rrule_env_transform(rrij, Zi, Zj, Rs, Zs, dB_cyl, basis.cutoff) -# # assemble the forces -# F[:, i] += dB_drrij -# F[:, j] -= dB_drrij -# for n = 1:length(Js) +# # assemble the forces +# F[:, i] += dB_drrij +# F[:, j] -= dB_drrij +# for n = 1:length(Js) # k = Js[n] # dv = dB_dRs[:, n] # F[:, k] -= dv -# F[:, i] += 0.5 * dv -# F[:, j] += 0.5 * dv +# F[:, i] += 0.5 * dv +# F[:, j] += 0.5 * dv # end # end -# return F # TODO: this is probably the wrong format +# return F # TODO: this is probably the wrong format # end -# function virial(basis::ACEBondPotentialBasis, at::Atoms{T}) where {T} +# function virial(basis::ACEBondPotentialBasis, at::Atoms{T}) where {T} # V = zeros(SMatrix{3, 3, T}, length(basis)) # for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis) # Zi, Zj = at.Z[i], at.Z[j] -# # find the right ace model +# # find the right ace model # ace = _get_model(basis, Zi, Zj) # # transform the euclidean to cylindrical coordinates # env = env_transform(rrij, Zi, Zj, Rs, Zs, basis.cutoff) -# # evaluate +# # evaluate # dB_cyl = evaluate_d(ace, env) -# # transform back? +# # transform back? # dB_drrij, dB_dRs = rrule_env_transform(rrij, Zi, Zj, Rs, Zs, dB_cyl, basis.cutoff) # # assemble the virials # for iB = 1:length(basis) @@ -214,7 +214,7 @@ end # for q = 1:length(Rs) # V[iB] -= dB_dRs[iB, q] * Rs[q]' # end -# end +# end # end -# return V # TODO: double-check the format -# end \ No newline at end of file +# return V # TODO: double-check the format +# end diff --git a/src/bonds/bselectors.jl b/src/bonds/bselectors.jl index b60d500..600f81e 100644 --- a/src/bonds/bselectors.jl +++ b/src/bonds/bselectors.jl @@ -1,122 +1,124 @@ module BondSelectors -import ACEfrictionCore: AbstractSparseBasis, maxorder, Prodb, Onepb, OneParticleBasis, - degree, CategorySparseBasis, SparseBasis -import Base: filter +import ACEfrictionCore: AbstractSparseBasis, maxorder, Prodb, Onepb, OneParticleBasis, + degree, CategorySparseBasis, SparseBasis +import Base: filter struct SparseCylindricalBondBasis <: AbstractSparseBasis - maxorder::Int - weight::Dict - maxlevels::Dict - p::Float64 - besym::Symbol - bondsym::Symbol - envsym::Symbol - ksym::Symbol - lsym::Symbol - weight_cat::Dict + maxorder::Int + weight::AbstractDict + maxlevels::AbstractDict + p::Float64 + besym::Symbol + bondsym::Symbol + envsym::Symbol + ksym::Symbol + lsym::Symbol + weight_cat::AbstractDict end -@noinline function SparseCylindricalBondBasis(; - besym = :be, bondsym = :bond, envsym = :env, ksym = :k, - lsym = :l, - maxorder = nothing, - p = 1, - weight = Dict{Symbol, Float64}(), - default_maxdeg = nothing, - maxlevels = nothing, - weight_cat = Dict(bondsym => 1.0, envsym => 1.0), - ) - @show maxorder, default_maxdeg - if (default_maxdeg != nothing) && (maxlevels == nothing) - return SparseCylindricalBondBasis(maxorder, weight, - Dict("default" => default_maxdeg), - p, besym, bondsym, envsym, ksym, lsym, weight_cat) - elseif (default_maxdeg == nothing) && (maxlevels != nothing) - return SparseCylindricalBondBasis(maxorder, weight, maxlevels, - p, besym, bondsym, envsym, ksym, lsym, weight_cat) - else - @error """Either both or neither optional arguments `maxlevels` and - `default_maxdeg` were provided. To avoid ambiguity ensure that - exactly one of these arguments is provided.""" - end +@noinline function SparseCylindricalBondBasis(; + besym=:be, bondsym=:bond, envsym=:env, ksym=:k, + lsym=:l, + maxorder=nothing, + p=1, + weight=Dict{Symbol,Float64}(), + default_maxdeg=nothing, + maxlevels=nothing, + weight_cat=Dict(bondsym => 1.0, envsym => 1.0), +) + @show maxorder, default_maxdeg + if (default_maxdeg != nothing) && (maxlevels == nothing) + return SparseCylindricalBondBasis(maxorder, weight, + Dict("default" => default_maxdeg), + p, besym, bondsym, envsym, ksym, lsym, weight_cat) + elseif (default_maxdeg == nothing) && (maxlevels != nothing) + return SparseCylindricalBondBasis(maxorder, weight, maxlevels, + p, besym, bondsym, envsym, ksym, lsym, weight_cat) + else + @error """Either both or neither optional arguments `maxlevels` and + `default_maxdeg` were provided. To avoid ambiguity ensure that + exactly one of these arguments is provided.""" + end end function Base.filter(b::Onepb, Bsel::SparseCylindricalBondBasis, basis::OneParticleBasis) - d = degree(b, basis) - k = b[Bsel.ksym] - return (k == 1 && b[Bsel.besym] == Bsel.envsym) || - (d == k-1 && b[Bsel.besym] == Bsel.bondsym) + d = degree(b, basis) + k = b[Bsel.ksym] + return (k == 1 && b[Bsel.besym] == Bsel.envsym) || + (d == k - 1 && b[Bsel.besym] == Bsel.bondsym) end function Base.filter(bb::Prodb, Bsel::SparseCylindricalBondBasis, basis::OneParticleBasis) - if length(bb) == 0; return true; end - has1bond = count((b[Bsel.besym] == Bsel.bondsym) for b in bb) == 1 - isinvariant = sum( b[Bsel.lsym] for b in bb ) == 0 - return has1bond && isinvariant + if length(bb) == 0 + return true + end + has1bond = count((b[Bsel.besym] == Bsel.bondsym) for b in bb) == 1 + isinvariant = sum(b[Bsel.lsym] for b in bb) == 0 + return has1bond && isinvariant end -# maxorder and maxlevel are inherited from the abstract interface +# maxorder and maxlevel are inherited from the abstract interface -level(b::Union{Prodb, Onepb}, Bsel::SparseCylindricalBondBasis, basis::OneParticleBasis) = - cat_weighted_degree(b, Bsel, basis) +level(b::Union{Prodb,Onepb}, Bsel::SparseCylindricalBondBasis, basis::OneParticleBasis) = + cat_weighted_degree(b, Bsel, basis) # Category-weighted degree function cat_weighted_degree(b::Onepb, Bsel::SparseCylindricalBondBasis, basis::OneParticleBasis) = - degree(b, basis, Bsel.weight) * Bsel.weight_cat[getproperty(b, Bsel.isym)] + degree(b, basis, Bsel.weight) * Bsel.weight_cat[getproperty(b, Bsel.isym)] cat_weighted_degree(bb::Prodb, Bsel::SparseCylindricalBondBasis, basis::OneParticleBasis) = ( - length(bb) == 0 ? 0.0 - : norm(cat_weighted_degree.(bb, Ref(Bsel), Ref(basis)), Bsel.p) - ) + length(bb) == 0 ? 0.0 + : norm(cat_weighted_degree.(bb, Ref(Bsel), Ref(basis)), Bsel.p) +) """ -Constructors of this alias of ACEfrictionCore.CategorySparseBasis can be used to conveniently create -basis selectors for ACE bond bases that are defined on ellipsoid-shaped bond environments -implmeneted as `EllipsoidCutoff` in the sub-module `ACEbonds.BondCutoffs``. +Constructors of this alias of ACEfrictionCore.CategorySparseBasis can be used to conveniently create +basis selectors for ACE bond bases that are defined on ellipsoid-shaped bond environments +implmeneted as `EllipsoidCutoff` in the sub-module `ACEbonds.BondCutoffs``. """ const EllipsoidBondBasis = CategorySparseBasis -function EllipsoidBondBasis(Bsel::SparseBasis; - isym=:mube, bond_weight = 1.0, species =[:X], - species_minorder_dict = Dict{Symbol,Int64}(), species_maxorder_dict = Dict{Symbol,Int64}(), - species_weight_cat = Dict( s => 1.0 for s in species)) - return CategorySparseBasis(isym, cat([:bond],species,dims=1); - maxorder = maxorder(Bsel), - p = Bsel.p, - weight = Bsel.weight, - maxlevels = Bsel.maxlevels, - minorder_dict = merge(Dict( :bond => 1), species_minorder_dict), - maxorder_dict = merge(Dict( :bond => 1), species_maxorder_dict), - weight_cat = merge(Dict(:bond => bond_weight), species_weight_cat) - ) +function EllipsoidBondBasis(Bsel::SparseBasis; + isym=:mube, bond_weight=1.0, species=[:X], + species_minorder_dict=Dict{Symbol,Int64}(), species_maxorder_dict=Dict{Symbol,Int64}(), + species_weight_cat=Dict(s => 1.0 for s in species)) + return CategorySparseBasis(isym, cat([:bond], species, dims=1); + maxorder=maxorder(Bsel), + p=Bsel.p, + weight=Bsel.weight, + maxlevels=Bsel.maxlevels, + minorder_dict=merge(Dict(:bond => 1), species_minorder_dict), + maxorder_dict=merge(Dict(:bond => 1), species_maxorder_dict), + weight_cat=merge(Dict(:bond => bond_weight), species_weight_cat) + ) end -function EllipsoidBondBasis(; - maxorder::Integer = nothing, - p = 1, - weight = Dict(:l => 1.0, :n => 1.0), - default_maxdeg = nothing, - maxlevels::Dict{Any, Float64} = nothing, - isym=:mube, bond_weight = 1.0, species =[:X], - species_minorder_dict = Dict{Any, Float64}(), - species_maxorder_dict = Dict{Any, Float64}(), - species_weight_cat = Dict(c => 1.0 for c in species), - ) - Bsel = SparseBasis(; maxorder = maxorder, - p = p, - weight = weight, - default_maxdeg = default_maxdeg, - maxlevels = maxlevels ) - return EllipsoidBondBasis(Bsel; - isym=isym, bond_weight = bond_weight, species = species, - species_minorder_dict = species_minorder_dict, species_maxorder_dict = species_maxorder_dict, - species_weight_cat = species_weight_cat) +function EllipsoidBondBasis(; + maxorder::Integer=nothing, + p=1, + weight=Dict(:l => 1.0, :n => 1.0), + default_maxdeg=nothing, + maxlevels::AbstractDict{Any,Float64}=nothing, + isym=:mube, bond_weight=1.0, species=[:X], + species_minorder_dict=Dict{Any,Float64}(), + species_maxorder_dict=Dict{Any,Float64}(), + species_weight_cat=Dict(c => 1.0 for c in species), +) + Bsel = SparseBasis(; maxorder=maxorder, + p=p, + weight=weight, + default_maxdeg=default_maxdeg, + maxlevels=maxlevels) + return EllipsoidBondBasis(Bsel; + isym=isym, bond_weight=bond_weight, species=species, + species_minorder_dict=species_minorder_dict, species_maxorder_dict=species_maxorder_dict, + species_weight_cat=species_weight_cat) end - -end \ No newline at end of file + +end diff --git a/src/bonds/iterator.jl b/src/bonds/iterator.jl index 7d23a71..352aafb 100644 --- a/src/bonds/iterator.jl +++ b/src/bonds/iterator.jl @@ -1,128 +1,133 @@ using JuLIP, StaticArrays, LinearAlgebra -import ACEfrictionCore: State, filter +import ACEfrictionCore: State, filter using JuLIP.Potentials: neigsz using JuLIP: Atoms # using ACEfrictionCore: BondEnvelope, filter, State, CylindricalBondEnvelope import ACEfrictionCore.ACEbonds.BondCutoffs: env_cutoff import ACEfrictionCore.ACEbonds.BondCutoffs: AbstractBondCutoff, env_filter, EllipsoidCutoff -_msort(z1,z2) = (z1<=z2 ? (z1,z2) : (z2,z1)) #TODO: this is hack. Need to either not use it here or define it once across all packages. +_msort(z1, z2) = (z1 <= z2 ? (z1, z2) : (z2, z1)) #TODO: this is hack. Need to either not use it here or define it once across all packages. #env_cutoff(cutoff::EllipsoidCutoff) = max(cutoff.rcutbond*.5 + cutoff.zcutenv, sqrt((cutoff.rcutbond*.5)^2+ cutoff.rcutenv^2)) -env_cutoff(cutoffs::Dict{Tuple{AtomicNumber,AtomicNumber},CUTOFF}) where {CUTOFF<:AbstractBondCutoff} = maximum(env_cutoff(c) for c in values(cutoffs)) +env_cutoff(cutoffs::AbstractDict{Tuple{AtomicNumber,AtomicNumber},CUTOFF}) where {CUTOFF<:AbstractBondCutoff} = maximum(env_cutoff(c) for c in values(cutoffs)) -bonds(at::Atoms, env::AbstractBondCutoff, args...) = - bonds( at, env.rcutbond, env_cutoff(env), - (r, z) -> env_filter(r, z, env), args...) +bonds(at::Atoms, env::AbstractBondCutoff, args...) = + bonds(at, env.rcutbond, env_cutoff(env), + (r, z) -> env_filter(r, z, env), args...) -# function bonds(at::Atoms, rcutbond, rcutenv, env_filter; subset=nothing) -# return (subset === nothing ? bonds(at::Atoms, rcutbond, rcutenv, env_filter) +# function bonds(at::Atoms, rcutbond, rcutenv, env_filter; subset=nothing) +# return (subset === nothing ? bonds(at::Atoms, rcutbond, rcutenv, env_filter) # : bonds(at::Atoms, rcutbond, rcutenv, env_filter,subset)) # end # TODO: make this type-stable -struct BondsIterator - at - nlist_bond - nlist_env - filter -end +struct BondsIterator + at + nlist_bond + nlist_env + filter +end """ -* rcutbond: include all bonds (i,j) such that rij <= rcutbond -* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` +* rcutbond: include all bonds (i,j) such that rij <= rcutbond +* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` * `filter` : `filter(X) == true` if particle `X` is to be included; `false` if to be discarded from the environment """ -function bonds(at::Atoms, rcutbond, rcutenv, filter) - nlist_bond = neighbourlist(at, rcutbond; recompute=true, storelist=false) - nlist_env = neighbourlist(at, rcutenv; recompute=true, storelist=false) - return BondsIterator(at, nlist_bond, nlist_env, filter) +function bonds(at::Atoms, rcutbond, rcutenv, filter) + nlist_bond = neighbourlist(at, rcutbond; recompute=true, storelist=false) + nlist_env = neighbourlist(at, rcutenv; recompute=true, storelist=false) + return BondsIterator(at, nlist_bond, nlist_env, filter) end -function Base.iterate(iter::BondsIterator, state=(1,0)) - i, q = state - # store temporary arrays for those... - Js, Rs = neigs(iter.nlist_bond, i) - - # nothing left to do - if i >= length(iter.at) && q >= length(Js) - return nothing - end - - # increment: - # case 1: we haven't yet exhausted the current neighbours. - # just increment the q index pointing into Js - if q < length(Js) - q += 1 - # here we could build in a rule to skip any pair for which we don't - # want to do the computation. - - # case 2: if i < length(at) but q >= length(Js) then we need to - # increment the i index and get the new neighbours - elseif i < length(iter.at) - i += 1 - Js, Rs = neigs(iter.nlist_bond, i) - q = 1 - else - return nothing - end - - j = Js[q] # index of neighbour (in central cell) - rr0 = rrij = Rs[q] # position of neighbour (in shifted cell) relative to i - # ssj = Rs[q] - iter.at.X[j] # shift of atom j into shifted cell - - # now we construct the environment - Js_e, Rs_e, Zs_e = _get_bond_env(iter, i, j, rrij) - - return (i, j, rrij, Js_e, Rs_e, Zs_e), (i, q) +function Base.iterate(iter::BondsIterator, state=(1, 0)) + i, q = state + # store temporary arrays for those... + Js, Rs = neigs(iter.nlist_bond, i) + + # nothing left to do + if i >= length(iter.at) && q >= length(Js) + return nothing + end + + # increment: + # case 1: we haven't yet exhausted the current neighbours. + # just increment the q index pointing into Js + if q < length(Js) + q += 1 + # here we could build in a rule to skip any pair for which we don't + # want to do the computation. + + # case 2: if i < length(at) but q >= length(Js) then we need to + # increment the i index and get the new neighbours + elseif i < length(iter.at) + i += 1 + Js, Rs = neigs(iter.nlist_bond, i) + q = 1 + else + return nothing + end + + j = Js[q] # index of neighbour (in central cell) + rr0 = rrij = Rs[q] # position of neighbour (in shifted cell) relative to i + # ssj = Rs[q] - iter.at.X[j] # shift of atom j into shifted cell + + # now we construct the environment + Js_e, Rs_e, Zs_e = _get_bond_env(iter, i, j, rrij) + + return (i, j, rrij, Js_e, Rs_e, Zs_e), (i, q) end function _get_bond_env(iter::BondsIterator, i, j, rrij) - # TODO: store temporary arrays - Js_i, Rs_i, Zs_i = neigsz(iter.nlist_env, iter.at, i) - - rri = iter.at.X[i] - rrmid = rri + 0.5 * rrij - Js = Int[]; sizehint!(Js, length(Js_i) ÷ 4) - Rs = typeof(rrij)[]; sizehint!(Rs, length(Js_i) ÷ 4) - Zs = AtomicNumber[]; sizehint!(Zs, length(Js_i) ÷ 4) - - ŝ = rrij/norm(rrij) - - # find the bond and remember it; - # TODO: this could now be integrated into the second loop - q_bond = 0 - for (q, rrq) in enumerate(Rs_i) - # rr = rrq + rri - rrmid - if rrq ≈ rrij # TODO: replace this with checking for j and shift! - @assert Js_i[q] == j - q_bond = q - break - end - end - if q_bond == 0 - error("the central bond neigbour atom j was not found") - end - - # now add the environment - for (q, rrq) in enumerate(Rs_i) - # skip the central bond - if q == q_bond; continue; end - # add the rest provided they fall within the provided filter - rr = rrq + rri - rrmid - z = dot(rr, ŝ) - r = norm(rr - z * ŝ) - if iter.filter(r, z) - push!(Js, Js_i[q]) - push!(Rs, rr) - push!(Zs, Zs_i[q]) - end - end - - return Js, Rs, Zs + # TODO: store temporary arrays + Js_i, Rs_i, Zs_i = neigsz(iter.nlist_env, iter.at, i) + + rri = iter.at.X[i] + rrmid = rri + 0.5 * rrij + Js = Int[] + sizehint!(Js, length(Js_i) ÷ 4) + Rs = typeof(rrij)[] + sizehint!(Rs, length(Js_i) ÷ 4) + Zs = AtomicNumber[] + sizehint!(Zs, length(Js_i) ÷ 4) + + ŝ = rrij / norm(rrij) + + # find the bond and remember it; + # TODO: this could now be integrated into the second loop + q_bond = 0 + for (q, rrq) in enumerate(Rs_i) + # rr = rrq + rri - rrmid + if rrq ≈ rrij # TODO: replace this with checking for j and shift! + @assert Js_i[q] == j + q_bond = q + break + end + end + if q_bond == 0 + error("the central bond neigbour atom j was not found") + end + + # now add the environment + for (q, rrq) in enumerate(Rs_i) + # skip the central bond + if q == q_bond + continue + end + # add the rest provided they fall within the provided filter + rr = rrq + rri - rrmid + z = dot(rr, ŝ) + r = norm(rr - z * ŝ) + if iter.filter(r, z) + push!(Js, Js_i[q]) + push!(Rs, rr) + push!(Zs, Zs_i[q]) + end + end + + return Js, Rs, Zs end @@ -134,293 +139,301 @@ end struct FilteredBondsIterator - at - nlist_bond - nlist_env - env_filter - subset -end + at + nlist_bond + nlist_env + env_filter + subset +end """ -* rcutbond: include all bonds (i,j) such that rij <= rcutbond -* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` +* rcutbond: include all bonds (i,j) such that rij <= rcutbond +* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` * `env_filter` : `env_filter(X) == true` if particle `X` is to be included; `false` if to be discarded from the environment -* `subset` : can either be of type Array{<:Int} in which case the bond iterator iterates only over bonds between atom pairs where the indices of both atoms are contained in indsf. +* `subset` : can either be of type Array{<:Int} in which case the bond iterator iterates only over bonds between atom pairs where the indices of both atoms are contained in indsf. Alternatively, indsf can also be of the form of a filter function `atom_filter(i::Int,at::AbstractAtoms)::Bool`, that returns `true` if bonds to the ith atom in the configuration `at` are to be included in the iterator, and `false`` otherwise. Consequently, the iterator only iterates over bonds between atom pairs where both atoms satisfy the filter criterion. """ bonds(at::Atoms, rcutbond, rcutenv, env_filter, subset) = FilteredBondsIterator(at, rcutbond, rcutenv, env_filter, subset) -bonds(at::Atoms, cutoff::AbstractBondCutoff, filter=(_,_)->true) = FilteredBondsIterator( at, cutoff.rcutbond, - env_cutoff(cutoff) , - (r, z) -> env_filter(r, z, cutoff), filter ) +bonds(at::Atoms, cutoff::AbstractBondCutoff, filter=(_, _) -> true) = FilteredBondsIterator(at, cutoff.rcutbond, + env_cutoff(cutoff), + (r, z) -> env_filter(r, z, cutoff), filter) """ -* rcutbond: include all bonds (i,j) such that rij <= rcutbond -* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` +* rcutbond: include all bonds (i,j) such that rij <= rcutbond +* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` * `env_filter` : `env_filter(X) == true` if particle `X` is to be included; `false` if to be discarded from the environment """ -function FilteredBondsIterator(at::Atoms, rcutbond::Real, rcutenv::Real, env_filter, subset::Array{<:Int}) - nlist_bond = neighbourlist(at, rcutbond; recompute=true, storelist=false) - nlist_env = neighbourlist(at, rcutenv; recompute=true, storelist=false) - return FilteredBondsIterator(at, nlist_bond, nlist_env, env_filter, subset) +function FilteredBondsIterator(at::Atoms, rcutbond::Real, rcutenv::Real, env_filter, subset::Array{<:Int}) + nlist_bond = neighbourlist(at, rcutbond; recompute=true, storelist=false) + nlist_env = neighbourlist(at, rcutenv; recompute=true, storelist=false) + return FilteredBondsIterator(at, nlist_bond, nlist_env, env_filter, subset) end -function FilteredBondsIterator(at::Atoms, rcutbond::Real, rcutenv::Real, env_filter, filter) - subset = findall(i->filter(i,at), 1:length(at) ) +function FilteredBondsIterator(at::Atoms, rcutbond::Real, rcutenv::Real, env_filter, filter) + subset = findall(i -> filter(i, at), 1:length(at)) #@show inds - return FilteredBondsIterator(at, rcutbond, rcutenv, env_filter, subset) + return FilteredBondsIterator(at, rcutbond, rcutenv, env_filter, subset) end function increment(iter::FilteredBondsIterator, state) ic, ib, Js, Rs = state ib = ib + 1 # increase bond index - if ib > length(Js) # already visited/iterated over all atoms in environment ? - ic = ic + 1 # increase index of center atom - if ic > length(iter.subset) # all relevant center atoms already visited? - return (nothing, ib, Js, Rs) # if yes, done! + if ib > length(Js) # already visited/iterated over all atoms in environment ? + ic = ic + 1 # increase index of center atom + if ic > length(iter.subset) # all relevant center atoms already visited? + return (nothing, ib, Js, Rs) # if yes, done! else ib = 1 # if no start at first atom in next environment Js, Rs = neigs(iter.nlist_bond, iter.subset[ic]) end - end + end return (ic, ib, Js, Rs) end function Base.iterate(iter::FilteredBondsIterator) - # if none of the atoms satisfy the filter criterion, there is nothing to iterate over - if length(iter.subset) == 0 - return nothing - else - Js, Rs = neigs(iter.nlist_bond, iter.subset[1]) - state = (1,0,Js,Rs) - return iterate(iter, state) - end + # if none of the atoms satisfy the filter criterion, there is nothing to iterate over + if length(iter.subset) == 0 + return nothing + else + Js, Rs = neigs(iter.nlist_bond, iter.subset[1]) + state = (1, 0, Js, Rs) + return iterate(iter, state) + end end function Base.iterate(iter::FilteredBondsIterator, state) - ic, ib, Js, Rs = state - - # Check whether s must be incremented (jumpt to next centre atom) or nothing left to iterate over - if ic > length(iter.subset) # nothing left to do - return nothing - end - #println("Before while") - #@show Js - while(true) + ic, ib, Js, Rs = state + + # Check whether s must be incremented (jumpt to next centre atom) or nothing left to iterate over + if ic > length(iter.subset) # nothing left to do + return nothing + end + #println("Before while") + #@show Js + while (true) (ic, ib, Js, Rs) = increment(iter, (ic, ib, Js, Rs)) if isnothing(ic) return nothing elseif !isempty(Js) && Js[ib] in iter.subset # here we could add a finer filter criterion, e.g. iter.fiter(iter.subset[ic], Js[ib], iter.at ) break end - end - i = iter.subset[ic] - j = Js[ib] # index of neighbour (in central cell) - rrij = Rs[ib] # position of neighbour (in shifted cell) relative to i - # ssj = Rs[q] - iter.at.X[j] # shift of atom j into shifted cell - # @show (i,j) - # now we construct the environment - Js_e, Rs_e, Zs_e = _get_bond_env(iter, i, j, rrij) - - return (i, j, rrij, Js_e, Rs_e, Zs_e), (ic, ib, Js, Rs) + end + i = iter.subset[ic] + j = Js[ib] # index of neighbour (in central cell) + rrij = Rs[ib] # position of neighbour (in shifted cell) relative to i + # ssj = Rs[q] - iter.at.X[j] # shift of atom j into shifted cell + # @show (i,j) + # now we construct the environment + Js_e, Rs_e, Zs_e = _get_bond_env(iter, i, j, rrij) + + return (i, j, rrij, Js_e, Rs_e, Zs_e), (ic, ib, Js, Rs) end function _get_bond_env(iter::FilteredBondsIterator, i, j, rrij) - # TODO: store temporary arrays - Js_i, Rs_i, Zs_i = neigsz(iter.nlist_env, iter.at, i) - - rri = iter.at.X[i] - rrmid = rri + 0.5 * rrij - Js = Int[]; sizehint!(Js, length(Js_i) ÷ 4) - Rs = typeof(rrij)[]; sizehint!(Rs, length(Js_i) ÷ 4) - Zs = AtomicNumber[]; sizehint!(Zs, length(Js_i) ÷ 4) - - ŝ = rrij/norm(rrij) - - # find the bond and remember it; - # TODO: this could now be integrated into the second loop - q_bond = 0 - for (q, rrq) in enumerate(Rs_i) - # rr = rrq + rri - rrmid - if rrq ≈ rrij # TODO: replace this with checking for j and shift! - @assert Js_i[q] == j - q_bond = q - break - end - end - if q_bond == 0 - error("the central bond neigbour atom j was not found") - end - - # now add the environment - for (q, rrq) in enumerate(Rs_i) - # skip the central bond - if q == q_bond; continue; end - # add the rest provided they fall within the provided env_filter - rr = rrq + rri - rrmid - z = dot(rr, ŝ) - r = norm(rr - z * ŝ) - if iter.env_filter(r, z) #TODO: by modifying the env_filter function we could allow for species pair-dependent Ellipsoid cutoffs. - push!(Js, Js_i[q]) - push!(Rs, rr) - push!(Zs, Zs_i[q]) - end - end - - return Js, Rs, Zs + # TODO: store temporary arrays + Js_i, Rs_i, Zs_i = neigsz(iter.nlist_env, iter.at, i) + + rri = iter.at.X[i] + rrmid = rri + 0.5 * rrij + Js = Int[] + sizehint!(Js, length(Js_i) ÷ 4) + Rs = typeof(rrij)[] + sizehint!(Rs, length(Js_i) ÷ 4) + Zs = AtomicNumber[] + sizehint!(Zs, length(Js_i) ÷ 4) + + ŝ = rrij / norm(rrij) + + # find the bond and remember it; + # TODO: this could now be integrated into the second loop + q_bond = 0 + for (q, rrq) in enumerate(Rs_i) + # rr = rrq + rri - rrmid + if rrq ≈ rrij # TODO: replace this with checking for j and shift! + @assert Js_i[q] == j + q_bond = q + break + end + end + if q_bond == 0 + error("the central bond neigbour atom j was not found") + end + + # now add the environment + for (q, rrq) in enumerate(Rs_i) + # skip the central bond + if q == q_bond + continue + end + # add the rest provided they fall within the provided env_filter + rr = rrq + rri - rrmid + z = dot(rr, ŝ) + r = norm(rr - z * ŝ) + if iter.env_filter(r, z) #TODO: by modifying the env_filter function we could allow for species pair-dependent Ellipsoid cutoffs. + push!(Js, Js_i[q]) + push!(Rs, rr) + push!(Zs, Zs_i[q]) + end + end + + return Js, Rs, Zs end struct FilteredBondsIteratorVarCutoff - at - nlist_bond - nlist_env - subset - cutoffs -end + at + nlist_bond + nlist_env + subset + cutoffs +end """ -* rcutbond: include all bonds (i,j) such that rij <= rcutbond -* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` +* rcutbond: include all bonds (i,j) such that rij <= rcutbond +* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` * `env_filter` : `env_filter(r,z,zzi,zzj) == true` if particle `X` is to be included; `false` if to be discarded from the environment -* `subset` : can either be of type Array{<:Int} in which case the bond iterator iterates only over bonds between atom pairs where the indices of both atoms are contained in indsf. +* `subset` : can either be of type Array{<:Int} in which case the bond iterator iterates only over bonds between atom pairs where the indices of both atoms are contained in indsf. Alternatively, indsf can also be of the form of a filter function `atom_filter(i::Int,at::AbstractAtoms)::Bool`, that returns `true` if bonds to the ith atom in the configuration `at` are to be included in the iterator, and `false`` otherwise. Consequently, the iterator only iterates over bonds between atom pairs where both atoms satisfy the filter criterion. """ -function bonds(at::Atoms, cutoffs::Dict{Tuple{AtomicNumber,AtomicNumber},CUTOFF}, subset::Array{<:Int}) where {CUTOFF<:AbstractBondCutoff} - rcutbond = maximum(cutoff.rcutbond for cutoff in values(cutoffs)) - rcutenv = env_cutoff(cutoffs) - return FilteredBondsIteratorVarCutoff(at, rcutbond, rcutenv,subset, cutoffs) +function bonds(at::Atoms, cutoffs::AbstractDict{Tuple{AtomicNumber,AtomicNumber},CUTOFF}, subset::Array{<:Int}) where {CUTOFF<:AbstractBondCutoff} + rcutbond = maximum(cutoff.rcutbond for cutoff in values(cutoffs)) + rcutenv = env_cutoff(cutoffs) + return FilteredBondsIteratorVarCutoff(at, rcutbond, rcutenv, subset, cutoffs) end -function bonds(at::Atoms, cutoffs::Dict{Tuple{AtomicNumber,AtomicNumber},CUTOFF}, filter= _->true) where {CUTOFF<:AbstractBondCutoff} - subset = findall(i->filter(i,at), 1:length(at) ) - return bonds(at, cutoffs, subset) +function bonds(at::Atoms, cutoffs::AbstractDict{Tuple{AtomicNumber,AtomicNumber},CUTOFF}, filter=_ -> true) where {CUTOFF<:AbstractBondCutoff} + subset = findall(i -> filter(i, at), 1:length(at)) + return bonds(at, cutoffs, subset) end """ -* rcutbond: include all bonds (i,j) such that rij <= rcutbond -* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` +* rcutbond: include all bonds (i,j) such that rij <= rcutbond +* `rcutenv`: include all bond environment atoms k such that `|rk - mid| <= rcutenv` * `env_filter` : `env_filter(X) == true` if particle `X` is to be included; `false` if to be discarded from the environment """ -function FilteredBondsIteratorVarCutoff(at::Atoms, rcutbond::Real, rcutenv::Real, subset::Array{<:Int}, cutoffs) - nlist_bond = neighbourlist(at, rcutbond; recompute=true, storelist=false) - nlist_env = neighbourlist(at, rcutenv; recompute=true, storelist=false) - return FilteredBondsIteratorVarCutoff(at, nlist_bond, nlist_env, subset, cutoffs) +function FilteredBondsIteratorVarCutoff(at::Atoms, rcutbond::Real, rcutenv::Real, subset::Array{<:Int}, cutoffs) + nlist_bond = neighbourlist(at, rcutbond; recompute=true, storelist=false) + nlist_env = neighbourlist(at, rcutenv; recompute=true, storelist=false) + return FilteredBondsIteratorVarCutoff(at, nlist_bond, nlist_env, subset, cutoffs) end -function FilteredBondsIteratorVarCutoff(at::Atoms, rcutbond::Real, rcutenv::Real, env_filter, filter) - subset = findall(i->filter(i,at), 1:length(at) ) +function FilteredBondsIteratorVarCutoff(at::Atoms, rcutbond::Real, rcutenv::Real, env_filter, filter) + subset = findall(i -> filter(i, at), 1:length(at)) #@show inds - return FilteredBondsIteratorVarCutoff(at, rcutbond, rcutenv, env_filter, subset) + return FilteredBondsIteratorVarCutoff(at, rcutbond, rcutenv, env_filter, subset) end function increment(iter::FilteredBondsIteratorVarCutoff, state) ic, ib, Js, Rs = state ib = ib + 1 # increase bond index - if ib > length(Js) # already visited/iterated over all atoms in environment ? - ic = ic + 1 # increase index of center atom - if ic > length(iter.subset) # all relevant center atoms already visited? - return (nothing, ib, Js, Rs) # if yes, done! + if ib > length(Js) # already visited/iterated over all atoms in environment ? + ic = ic + 1 # increase index of center atom + if ic > length(iter.subset) # all relevant center atoms already visited? + return (nothing, ib, Js, Rs) # if yes, done! else ib = 1 # if no start at first atom in next environment Js, Rs = neigs(iter.nlist_bond, iter.subset[ic]) end - end + end return (ic, ib, Js, Rs) end function Base.iterate(iter::FilteredBondsIteratorVarCutoff) - # if none of the atoms satisfy the filter criterion, there is nothing to iterate over - if length(iter.subset) == 0 - return nothing - else - Js, Rs = neigs(iter.nlist_bond, iter.subset[1]) - state = (1,0,Js,Rs) - return iterate(iter, state) - end + # if none of the atoms satisfy the filter criterion, there is nothing to iterate over + if length(iter.subset) == 0 + return nothing + else + Js, Rs = neigs(iter.nlist_bond, iter.subset[1]) + state = (1, 0, Js, Rs) + return iterate(iter, state) + end end function Base.iterate(iter::FilteredBondsIteratorVarCutoff, state) - ic, ib, Js, Rs = state - Zs = iter.at.Z[Js] - # Check whether s must be incremented (jumpt to next centre atom) or nothing left to iterate over - if ic > length(iter.subset) # nothing left to do - return nothing - end - #println("Before while") - #@show Js - while(true) + ic, ib, Js, Rs = state + Zs = iter.at.Z[Js] + # Check whether s must be incremented (jumpt to next centre atom) or nothing left to iterate over + if ic > length(iter.subset) # nothing left to do + return nothing + end + #println("Before while") + #@show Js + while (true) (ic, ib, Js, Rs) = increment(iter, (ic, ib, Js, Rs)) if isnothing(ic) return nothing - elseif !isempty(Js) && Js[ib] in iter.subset && haskey(iter.cutoffs,_msort(iter.at.Z[iter.subset[ic]],iter.at.Z[Js[ib]])) && norm(Rs[ib]) < iter.cutoffs[_msort(iter.at.Z[iter.subset[ic]],iter.at.Z[Js[ib]])].rcutbond # here we could add a finer filter criterion, e.g. iter.fiter(iter.subset[ic], Js[ib], iter.at ) + elseif !isempty(Js) && Js[ib] in iter.subset && haskey(iter.cutoffs, _msort(iter.at.Z[iter.subset[ic]], iter.at.Z[Js[ib]])) && norm(Rs[ib]) < iter.cutoffs[_msort(iter.at.Z[iter.subset[ic]], iter.at.Z[Js[ib]])].rcutbond # here we could add a finer filter criterion, e.g. iter.fiter(iter.subset[ic], Js[ib], iter.at ) break end - end - i = iter.subset[ic] - j = Js[ib] # index of neighbour (in central cell) - rrij = Rs[ib] # position of neighbour (in shifted cell) relative to i - # ssj = Rs[q] - iter.at.X[j] # shift of atom j into shifted cell - # @show (i,j) - # now we construct the environment - Js_e, Rs_e, Zs_e = _get_bond_env(iter, i, j, rrij) - - return (i, j, rrij, Js_e, Rs_e, Zs_e), (ic, ib, Js, Rs) + end + i = iter.subset[ic] + j = Js[ib] # index of neighbour (in central cell) + rrij = Rs[ib] # position of neighbour (in shifted cell) relative to i + # ssj = Rs[q] - iter.at.X[j] # shift of atom j into shifted cell + # @show (i,j) + # now we construct the environment + Js_e, Rs_e, Zs_e = _get_bond_env(iter, i, j, rrij) + + return (i, j, rrij, Js_e, Rs_e, Zs_e), (ic, ib, Js, Rs) end function _get_bond_env(iter::FilteredBondsIteratorVarCutoff, i, j, rrij) - # TODO: store temporary arrays - Js_i, Rs_i, Zs_i = neigsz(iter.nlist_env, iter.at, i) - - rri = iter.at.X[i] - rrmid = rri + 0.5 * rrij - Js = Int[]; sizehint!(Js, length(Js_i) ÷ 4) - Rs = typeof(rrij)[]; sizehint!(Rs, length(Js_i) ÷ 4) - Zs = AtomicNumber[]; sizehint!(Zs, length(Js_i) ÷ 4) - - ŝ = rrij/norm(rrij) - - # find the bond and remember it; - # TODO: this could now be integrated into the second loop - q_bond = 0 - for (q, rrq) in enumerate(Rs_i) - # rr = rrq + rri - rrmid - if rrq ≈ rrij # TODO: replace this with checking for j and shift! - @assert Js_i[q] == j - q_bond = q - break - end - end - if q_bond == 0 - error("the central bond neigbour atom j was not found") - end - - # now add the environment - cutoff = iter.cutoffs[_msort(iter.at.Z[i], iter.at.Z[j])] - for (q, rrq) in enumerate(Rs_i) - # skip the central bond - if q == q_bond; continue; end - # add the rest provided they fall within the provided env_filter - rr = rrq + rri - rrmid - z = dot(rr, ŝ) - r = norm(rr - z * ŝ) - if env_filter(r, z, cutoff) #TODO: by modifying the env_filter function we could allow for species pair-dependent Ellipsoid cutoffs. - push!(Js, Js_i[q]) - push!(Rs, rr) - push!(Zs, Zs_i[q]) - end - end - - return Js, Rs, Zs -end - + # TODO: store temporary arrays + Js_i, Rs_i, Zs_i = neigsz(iter.nlist_env, iter.at, i) + + rri = iter.at.X[i] + rrmid = rri + 0.5 * rrij + Js = Int[] + sizehint!(Js, length(Js_i) ÷ 4) + Rs = typeof(rrij)[] + sizehint!(Rs, length(Js_i) ÷ 4) + Zs = AtomicNumber[] + sizehint!(Zs, length(Js_i) ÷ 4) + + ŝ = rrij / norm(rrij) + + # find the bond and remember it; + # TODO: this could now be integrated into the second loop + q_bond = 0 + for (q, rrq) in enumerate(Rs_i) + # rr = rrq + rri - rrmid + if rrq ≈ rrij # TODO: replace this with checking for j and shift! + @assert Js_i[q] == j + q_bond = q + break + end + end + if q_bond == 0 + error("the central bond neigbour atom j was not found") + end + + # now add the environment + cutoff = iter.cutoffs[_msort(iter.at.Z[i], iter.at.Z[j])] + for (q, rrq) in enumerate(Rs_i) + # skip the central bond + if q == q_bond + continue + end + # add the rest provided they fall within the provided env_filter + rr = rrq + rri - rrmid + z = dot(rr, ŝ) + r = norm(rr - z * ŝ) + if env_filter(r, z, cutoff) #TODO: by modifying the env_filter function we could allow for species pair-dependent Ellipsoid cutoffs. + push!(Js, Js_i[q]) + push!(Rs, rr) + push!(Zs, Zs_i[q]) + end + end + return Js, Rs, Zs +end diff --git a/src/bonds/utils.jl b/src/bonds/utils.jl index 64eb953..aa7569b 100644 --- a/src/bonds/utils.jl +++ b/src/bonds/utils.jl @@ -3,70 +3,69 @@ import ACEfrictionCore import ACEfrictionCore: polytransform export SymmetricEllipsoidBondBasis - # explicitly included all optional arguments for transparancy - function SymmetricEllipsoidBondBasis(ϕ::ACEfrictionCore.AbstractProperty; - maxorder::Integer = nothing, - p = 1, - weight = Dict(:l => 1.0, :n => 1.0), - default_maxdeg = nothing, - #maxlevels::Dict{Any, Float64} = nothing, - r0 = .4, - rin=.0, - trans = polytransform(2, r0), - pcut=2, - pin=2, - bondsymmetry=nothing, - kvargs...) # kvargs = additional optional arguments for EllipsoidBondBasis: i.e., species =[:X], isym=:mube, bond_weight = 1.0, species_minorder_dict = Dict{Any, Float64}(), species_maxorder_dict = Dict{Any, Float64}(), species_weight_cat = Dict(c => 1.0 for c in species), - Bsel = SparseBasis(; maxorder = maxorder, - p = p, - weight = weight, - default_maxdeg = default_maxdeg) - #maxlevels = maxlevels ) - return SymmetricEllipsoidBondBasis(ϕ, Bsel; r0=r0, rin=rin,trans=trans, pcut=pcut, pin=pin,bondsymmetry=bondsymmetry, kvargs...) - end +# explicitly included all optional arguments for transparancy +function SymmetricEllipsoidBondBasis(ϕ::ACEfrictionCore.AbstractProperty; + maxorder::Integer=nothing, + p=1, + weight=Dict(:l => 1.0, :n => 1.0), + default_maxdeg=nothing, + #maxlevels::AbstractDict{Any, Float64} = nothing, + r0=0.4, + rin=0.0, + trans=polytransform(2, r0), + pcut=2, + pin=2, + bondsymmetry=nothing, + kvargs...) # kvargs = additional optional arguments for EllipsoidBondBasis: i.e., species =[:X], isym=:mube, bond_weight = 1.0, species_minorder_dict = Dict{Any, Float64}(), species_maxorder_dict = Dict{Any, Float64}(), species_weight_cat = Dict(c => 1.0 for c in species), + Bsel = SparseBasis(; maxorder=maxorder, + p=p, + weight=weight, + default_maxdeg=default_maxdeg) + #maxlevels = maxlevels ) + return SymmetricEllipsoidBondBasis(ϕ, Bsel; r0=r0, rin=rin, trans=trans, pcut=pcut, pin=pin, bondsymmetry=bondsymmetry, kvargs...) +end + - - function SymmetricEllipsoidBondBasis(ϕ::ACEfrictionCore.AbstractProperty, Bsel::ACEfrictionCore.SparseBasis; - r0 = .4, - rin=.0, - trans = polytransform(2, r0), - pcut=2, - pin=2, - bondsymmetry=nothing, - species =[:X], - kvargs... - ) - if haskey(kvargs,:isym) - @assert kvargs[:isym] == :mube - end - @assert 0.0 < r0 < 1.0 - @assert 0.0 <= rin < 1.0 +function SymmetricEllipsoidBondBasis(ϕ::ACEfrictionCore.AbstractProperty, Bsel::ACEfrictionCore.SparseBasis; + r0=0.4, + rin=0.0, + trans=polytransform(2, r0), + pcut=2, + pin=2, + bondsymmetry=nothing, + species=[:X], + kvargs... +) + if haskey(kvargs, :isym) + @assert kvargs[:isym] == :mube + end + @assert 0.0 < r0 < 1.0 + @assert 0.0 <= rin < 1.0 - BondSelector = EllipsoidBondBasis( Bsel; species=species, kvargs...) - min_weight = minimum(values(BondSelector.weight_cat)) - maxdeg = Int(ceil(maximum(values(BondSelector.maxlevels)))) - RnYlm = ACEfrictionCore.Utils.RnYlm_1pbasis(; r0 = r0, - rin = rin, - trans = trans, - pcut = pcut, - pin = pin, - rcut= 1.0, - Bsel = Bsel, - maxdeg= maxdeg * max(1,Int(ceil(1/min_weight))) - ); - Bc = ACEfrictionCore.Categorical1pBasis(cat([:bond],species, dims=1); varsym = :mube, idxsym = :mube ) - B1p = Bc * RnYlm - return SymmetricEllipsoidBondBasis(ϕ, BondSelector, B1p; bondsymmetry=bondsymmetry) + BondSelector = EllipsoidBondBasis(Bsel; species=species, kvargs...) + min_weight = minimum(values(BondSelector.weight_cat)) + maxdeg = Int(ceil(maximum(values(BondSelector.maxlevels)))) + RnYlm = ACEfrictionCore.Utils.RnYlm_1pbasis(; r0=r0, + rin=rin, + trans=trans, + pcut=pcut, + pin=pin, + rcut=1.0, + Bsel=Bsel, + maxdeg=maxdeg * max(1, Int(ceil(1 / min_weight))) + ) + Bc = ACEfrictionCore.Categorical1pBasis(cat([:bond], species, dims=1); varsym=:mube, idxsym=:mube) + B1p = Bc * RnYlm + return SymmetricEllipsoidBondBasis(ϕ, BondSelector, B1p; bondsymmetry=bondsymmetry) end function SymmetricEllipsoidBondBasis(ϕ::ACEfrictionCore.AbstractProperty, BondSelector::EllipsoidBondBasis, B1p::ACEfrictionCore.Product1pBasis; bondsymmetry=nothing) - filterfun = _->true - if bondsymmetry == "Invariant" - filterfun = ACEfrictionCore.EvenL(:mube, [:bond]) - end - if bondsymmetry == "Covariant" - filterfun = x -> !(ACEfrictionCore.EvenL(:mube, [:bond])(x)) - end - return ACEfrictionCore.SymmetricBasis(ϕ, B1p, BondSelector; filterfun = filterfun) + filterfun = _ -> true + if bondsymmetry == "Invariant" + filterfun = ACEfrictionCore.EvenL(:mube, [:bond]) + end + if bondsymmetry == "Covariant" + filterfun = x -> !(ACEfrictionCore.EvenL(:mube, [:bond])(x)) + end + return ACEfrictionCore.SymmetricBasis(ϕ, B1p, BondSelector; filterfun=filterfun) end - diff --git a/src/chain.jl b/src/chain.jl index 6d0377f..0df1587 100644 --- a/src/chain.jl +++ b/src/chain.jl @@ -1,58 +1,57 @@ -abstract type AbstractSChain{TT} end +abstract type AbstractSChain{TT} end struct SChain{TT} <: AbstractSChain{TT} - F::TT + F::TT end -struct TypedChain{TT, IN, OUT} <: AbstractSChain{TT} - F::TT -end +struct TypedChain{TT,IN,OUT} <: AbstractSChain{TT} + F::TT +end -# construct a chain recursively -chain(F1, F2, args...) = chain( chain(F1, F2), args... ) -# for most arguments, just form a tuple -chain(F1, F2) = SChain( (F1, F2) ) -# if one of them is a chain already, then combine into a single long chain -chain(F1::SChain, F2) = SChain( tuple(F1.F..., F2) ) -chain(F1, F2::SChain) = SChain( tuple(F1, F2.F...) ) -chain(F1::SChain, F2::SChain) = chain( tuple(F1.F..., F2.F...) ) +# construct a chain recursively +chain(F1, F2, args...) = chain(chain(F1, F2), args...) +# for most arguments, just form a tuple +chain(F1, F2) = SChain((F1, F2)) +# if one of them is a chain already, then combine into a single long chain +chain(F1::SChain, F2) = SChain(tuple(F1.F..., F2)) +chain(F1, F2::SChain) = SChain(tuple(F1, F2.F...)) +chain(F1::SChain, F2::SChain) = chain(tuple(F1.F..., F2.F...)) Base.length(c::SChain) = length(c.F) @generated function evaluate(chain::AbstractSChain{TT}, X) where {TT} - LEN = length(TT.types) - code = Expr[] - push!(code, :(X_0 = X)) - for l = 1:LEN - push!(code, Meta.parse("F_$l = chain.F[$l]")) - push!(code, Meta.parse("X_$l = evaluate(F_$l, X_$(l-1))")) - push!(code, Meta.parse("release!(X_$(l-1))")) - end - push!(code, Meta.parse("return X_$LEN")) - return Expr(:block, code...) + LEN = length(TT.types) + code = Expr[] + push!(code, :(X_0 = X)) + for l = 1:LEN + push!(code, Meta.parse("F_$l = chain.F[$l]")) + push!(code, Meta.parse("X_$l = evaluate(F_$l, X_$(l-1))")) + push!(code, Meta.parse("release!(X_$(l-1))")) + end + push!(code, Meta.parse("return X_$LEN")) + return Expr(:block, code...) end -## +## -import Base: == +import Base: == -==(ch1::SChain, ch2::SChain) = - all( F1==F2 for (F1, F2) in zip(ch1.F, ch2.F) ) +==(ch1::SChain, ch2::SChain) = + all(F1 == F2 for (F1, F2) in zip(ch1.F, ch2.F)) write_dict(chain::SChain) = Dict( - "__id__" => "ACEfrictionCore_SChain", - "F" => write_dict.(chain.F) - ) - -read_dict(::Val{:ACEfrictionCore_SChain}, D::Dict) = - SChain(tuple( read_dict.(D["F"])... )) + "__id__" => "ACEfrictionCore_SChain", + "F" => write_dict.(chain.F) +) +read_dict(::Val{:ACEfrictionCore_SChain}, D::AbstractDict) = + SChain(tuple(read_dict.(D["F"])...)) -## ALTERNATIVE CHAIN IMPLEMENTATION - LINKS -abstract type ChainLink end +## ALTERNATIVE CHAIN IMPLEMENTATION - LINKS -next(::ChainLink) = nothing +abstract type ChainLink end -previous(::ChainLink) = nothing +next(::ChainLink) = nothing +previous(::ChainLink) = nothing diff --git a/src/discrete1pbasis.jl b/src/discrete1pbasis.jl index 0929c6a..b5a67ad 100644 --- a/src/discrete1pbasis.jl +++ b/src/discrete1pbasis.jl @@ -8,20 +8,20 @@ export Categorical1pBasis # ------------------------- -struct SList{N, T} - list::SVector{N, T} - - function SList{N, T}(list::SVector{N, T}) where {N, T} - if isabstracttype(T) - error("`SList` can only contain a single type") - end - return new(list) - end +struct SList{N,T} + list::SVector{N,T} + + function SList{N,T}(list::SVector{N,T}) where {N,T} + if isabstracttype(T) + error("`SList` can only contain a single type") + end + return new(list) + end end SList(list::AbstractArray) = SList(SVector(list...)) SList(args...) = SList(SVector(args...)) -SList(list::SVector{N, T}) where {N, T} = SList{N, T}(list) +SList(list::SVector{N,T}) where {N,T} = SList{N,T}(list) Base.length(list::SList) = length(list.list) Base.rand(list::SList) = rand(list.list) @@ -31,24 +31,24 @@ i2val(list::SList, i::Integer) = list.list[i] Base.iterate(list::SList, args...) = iterate(list.list, args...) function val2i(list::SList, val) - for j = 1:length(list) - if list.list[j] == val - return j - end - end - error("val = $val not found in this list") + for j = 1:length(list) + if list.list[j] == val + return j + end + end + error("val = $val not found in this list") end -write_dict(list::SList{N,T}) where {N, T} = - Dict( "__id__" => "ACEfrictionCore_SList", - "T" => write_dict(T), - "list" => list.list ) +write_dict(list::SList{N,T}) where {N,T} = + Dict("__id__" => "ACEfrictionCore_SList", + "T" => write_dict(T), + "list" => list.list) -function read_dict(::Val{:ACEfrictionCore_SList}, D::Dict) - list = D["list"] - T = read_dict(D["T"]) - svector = SVector{length(list), T}((T.(list))...) - return SList(svector) +function read_dict(::Val{:ACEfrictionCore_SList}, D::AbstractDict) + list = D["list"] + T = read_dict(D["T"]) + svector = SVector{length(list),T}((T.(list))...) + return SList(svector) end @@ -56,89 +56,89 @@ end # ------------------------- @doc raw""" -`Categorical1pBasis` : defines the discrete 1p basis -```math +`Categorical1pBasis` : defines the discrete 1p basis +```math \phi_q(u) = \delta(u - U_q), ``` -where ``U_q, q = 1, \dots, Q`` are finitely many possible values that the -variable ``u`` may take. Suppose, e.g., we allow the values `[:a, :b, :c]`, -then -```julia +where ``U_q, q = 1, \dots, Q`` are finitely many possible values that the +variable ``u`` may take. Suppose, e.g., we allow the values `[:a, :b, :c]`, +then +```julia P = Categorical1pBasis([:a, :b, :c]; varsym = :u, idxsym = :q) evaluate(P, State(u = :a)) # Bool[1, 0, 0] evaluate(P, State(u = :b)) # Bool[0, 1, 0] evaluate(P, State(u = :c)) # Bool[0, 0, 1] ``` -If we evaluate it with an unknown state we get an error: -```julia -evaluate(P, State(u = :x)) +If we evaluate it with an unknown state we get an error: +```julia +evaluate(P, State(u = :x)) # Error : val = x not found in this list ``` """ -struct Categorical1pBasis{VSYM, ISYM, LEN, T} <: Discrete1pBasis{Bool} - categories::SList{LEN, T} - label::String +struct Categorical1pBasis{VSYM,ISYM,LEN,T} <: Discrete1pBasis{Bool} + categories::SList{LEN,T} + label::String end -_varsym(::Categorical1pBasis{VSYM, ISYM}) where {VSYM, ISYM} = VSYM -_isym(::Categorical1pBasis{VSYM, ISYM}) where {VSYM, ISYM} = ISYM +_varsym(::Categorical1pBasis{VSYM,ISYM}) where {VSYM,ISYM} = VSYM +_isym(::Categorical1pBasis{VSYM,ISYM}) where {VSYM,ISYM} = ISYM _val(X, B::Categorical1pBasis) = getproperty(X, _varsym(B)) _idx(b, B::Categorical1pBasis) = getproperty(b, _isym(B)) Base.length(B::Categorical1pBasis) = length(B.categories) -Categorical1pBasis(categories::AbstractArray; - varsym::Symbol = nothing, idxsym::Symbol = nothing, - label = "C$(idxsym)") = - Categorical1pBasis(categories, varsym, idxsym, label) +Categorical1pBasis(categories::AbstractArray; + varsym::Symbol=nothing, idxsym::Symbol=nothing, + label="C$(idxsym)") = + Categorical1pBasis(categories, varsym, idxsym, label) -Categorical1pBasis(categories::AbstractArray, varsym::Symbol, isym::Symbol, label::String) = - Categorical1pBasis(SList(categories), varsym, isym, label) +Categorical1pBasis(categories::AbstractArray, varsym::Symbol, isym::Symbol, label::String) = + Categorical1pBasis(SList(categories), varsym, isym, label) -Categorical1pBasis(categories::SList{LEN, T}, varsym::Symbol, isym::Symbol, label::String) where {LEN, T} = - Categorical1pBasis{varsym, isym, LEN, T}(categories, label) +Categorical1pBasis(categories::SList{LEN,T}, varsym::Symbol, isym::Symbol, label::String) where {LEN,T} = + Categorical1pBasis{varsym,isym,LEN,T}(categories, label) -function evaluate(basis::Categorical1pBasis, X::AbstractState) - A = Vector{Bool}(undef, length(basis)) - return evaluate!(A, basis, X) +function evaluate(basis::Categorical1pBasis, X::AbstractState) + A = Vector{Bool}(undef, length(basis)) + return evaluate!(A, basis, X) end function ACEfrictionCore.evaluate!(A, basis::Categorical1pBasis, X::AbstractState) - fill!(A, false) - A[val2i(basis.categories, _val(X, basis))] = true - return A + fill!(A, false) + A[val2i(basis.categories, _val(X, basis))] = true + return A end -symbols(basis::Categorical1pBasis) = [ _isym(basis), ] +symbols(basis::Categorical1pBasis) = [_isym(basis),] -indexrange(basis::Categorical1pBasis) = Dict( _isym(basis) => basis.categories.list ) +indexrange(basis::Categorical1pBasis) = Dict(_isym(basis) => basis.categories.list) isadmissible(b, basis::Categorical1pBasis) = (_idx(b, basis) in basis.categories) -get_index(B::Categorical1pBasis, b) = val2i(B.categories, _idx(b, B) ) +get_index(B::Categorical1pBasis, b) = val2i(B.categories, _idx(b, B)) degree(b, basis::Categorical1pBasis, args...) = 0 Base.rand(basis::Categorical1pBasis) = rand(basis.list) function get_spec(basis::Categorical1pBasis, i) - return NamedTuple{(_isym(basis),)}((i2val(basis.categories, i),)) + return NamedTuple{(_isym(basis),)}((i2val(basis.categories, i),)) end -get_spec(basis::Categorical1pBasis) = [ get_spec(basis, i) for i = 1:length(basis) ] +get_spec(basis::Categorical1pBasis) = [get_spec(basis, i) for i = 1:length(basis)] -write_dict(B::Categorical1pBasis) = - Dict( "__id__" => "ACEfrictionCore_Categorical1pBasis", - "categories" => write_dict(B.categories), - "VSYM" => String(_varsym(B)), - "ISYM" => String(_isym(B)), - "label" => B.label) +write_dict(B::Categorical1pBasis) = + Dict("__id__" => "ACEfrictionCore_Categorical1pBasis", + "categories" => write_dict(B.categories), + "VSYM" => String(_varsym(B)), + "ISYM" => String(_isym(B)), + "label" => B.label) -read_dict(::Val{:ACEfrictionCore_Categorical1pBasis}, D::Dict) = - Categorical1pBasis( read_dict(D["categories"]), - Symbol(D["VSYM"]), Symbol(D["ISYM"]), - D["label"] ) +read_dict(::Val{:ACEfrictionCore_Categorical1pBasis}, D::AbstractDict) = + Categorical1pBasis(read_dict(D["categories"]), + Symbol(D["VSYM"]), Symbol(D["ISYM"]), + D["label"]) diff --git a/src/evaluator.jl b/src/evaluator.jl index 41b79fb..4f9cb26 100644 --- a/src/evaluator.jl +++ b/src/evaluator.jl @@ -3,75 +3,75 @@ """ `struct ProductEvaluator` : specifies a ProductEvaluator, which is basically defined -through a PIBasis and its coefficients. The n-correlations are evaluated directly +through a PIBasis and its coefficients. The n-correlations are evaluated directly via a naive product of the atomic base. """ -mutable struct ProductEvaluator{T, TPI <: PIBasis, REAL} - pibasis::TPI # AA basis from ACE papers - coeffs::Vector{T} # c̃ coefficients from ACE papers - real::REAL # the real operation stored in the SymmetricBasis +mutable struct ProductEvaluator{T,TPI<:PIBasis,REAL} + pibasis::TPI # AA basis from ACE papers + coeffs::Vector{T} # c̃ coefficients from ACE papers + real::REAL # the real operation stored in the SymmetricBasis end ==(V1::ProductEvaluator, V2::ProductEvaluator) = - (V1.pibasis == V2.pibasis) && (V1.coeffs == V2.coeffs) + (V1.pibasis == V2.pibasis) && (V1.coeffs == V2.coeffs) -# ----------- FIO +# ----------- FIO -write_dict(ev::ProductEvaluator) = Dict( "__id__" => "ACEfrictionCore_ProductEvaluator" ) +write_dict(ev::ProductEvaluator) = Dict("__id__" => "ACEfrictionCore_ProductEvaluator") -read_dict(::Val{:ACEfrictionCore_ProductEvaluator}, D::Dict, basis, c) = ProductEvaluator(basis, c) +read_dict(::Val{:ACEfrictionCore_ProductEvaluator}, D::AbstractDict, basis, c) = ProductEvaluator(basis, c) # ------------------------------------------------------------ # Initialisation and Parameter manipulation code # ------------------------------------------------------------ -ProductEvaluator(basis::SymmetricBasis, c::AbstractVector) = - ProductEvaluator(basis.pibasis, _get_eff_coeffs(basis, c), basis.real) +ProductEvaluator(basis::SymmetricBasis, c::AbstractVector) = + ProductEvaluator(basis.pibasis, _get_eff_coeffs(basis, c), basis.real) -# basic setter interface without any checks +# basic setter interface without any checks function set_params!(ev::ProductEvaluator, c̃::AbstractVector) - ev.coeffs[:] .= c̃ - return ev -end + ev.coeffs[:] .= c̃ + return ev +end # trivial setter when the parameters already come in the c̃ format (AA-basis) function set_params!(ev::ProductEvaluator, basis::PIBasis, c̃::AbstractVector) - @assert ev.pibasis === basis - set_params!(ev, c̃) + @assert ev.pibasis === basis + set_params!(ev, c̃) end -# if parameters come in c (B-basis) format then they first need to be converted +# if parameters come in c (B-basis) format then they first need to be converted # to c̃ format (AA basis) function set_params!(ev::ProductEvaluator, basis::SymmetricBasis, c::AbstractVector) - len_AA = length(ev.pibasis) - @assert len_AA == size(basis.A2Bmap, 2) - c̃ = _acquire_ctilde(basis, len_AA, c) - _get_eff_coeffs!(c̃, basis, c) - set_params!(ev, basis.pibasis, c̃) - release!(c̃) - return ev + len_AA = length(ev.pibasis) + @assert len_AA == size(basis.A2Bmap, 2) + c̃ = _acquire_ctilde(basis, len_AA, c) + _get_eff_coeffs!(c̃, basis, c) + set_params!(ev, basis.pibasis, c̃) + release!(c̃) + return ev end -_get_eff_coeffs!(c̃, basis::SymmetricBasis, c::AbstractVector) = - genmul!(c̃, transpose(basis.A2Bmap), c, *) +_get_eff_coeffs!(c̃, basis::SymmetricBasis, c::AbstractVector) = + genmul!(c̃, transpose(basis.A2Bmap), c, *) function _get_eff_coeffs(basis::SymmetricBasis, c::AbstractVector) - c̃ = _alloc_ctilde(basis,c) - return _get_eff_coeffs!(c̃, basis, c) + c̃ = _alloc_ctilde(basis, c) + return _get_eff_coeffs!(c̃, basis, c) end # TODO: we may need a second pool to allocate ctilde vectors... -struct _One <: Number +struct _One <: Number end -import Base: * +import Base: * *(x, ::_One) = x *(::_One, x) = x *(x::AbstractProperty, ::_One) = x @@ -81,31 +81,31 @@ import Base: * *(::ACEfrictionCore._One, x::ACEfrictionCore.XState) = x *(x::ACEfrictionCore.XState, ::ACEfrictionCore._One) = x -_acquire_ctilde(basis::SymmetricBasis, len_AA, c::AbstractVector{<: Number}) = - zeros(promote_type(eltype(basis.A2Bmap), eltype(c)), len_AA) +_acquire_ctilde(basis::SymmetricBasis, len_AA, c::AbstractVector{<:Number}) = + zeros(promote_type(eltype(basis.A2Bmap), eltype(c)), len_AA) + +_acquire_ctilde(basis::SymmetricBasis, len_AA, c::AbstractVector{<:SVector}) = + zeros(SVector{length(c[1]), + promote_type(eltype(basis.A2Bmap), eltype(c[1])) + }, len_AA) -_acquire_ctilde(basis::SymmetricBasis, len_AA, c::AbstractVector{<: SVector}) = - zeros(SVector{length(c[1]), - promote_type(eltype(basis.A2Bmap), eltype(c[1])) - } , len_AA ) +_alloc_ctilde(basis::SymmetricBasis, c::AbstractVector{<:SVector}) = + zeros(SVector{length(c[1]),eltype(basis.A2Bmap)}, size(basis.A2Bmap, 2)) -_alloc_ctilde(basis::SymmetricBasis, c::AbstractVector{<: SVector}) = - zeros(SVector{length(c[1]),eltype(basis.A2Bmap)}, size(basis.A2Bmap, 2)) - -_alloc_ctilde(basis::SymmetricBasis, c::AbstractVector{<: Number}) = - zeros(eltype(basis.A2Bmap), size(basis.A2Bmap, 2)) +_alloc_ctilde(basis::SymmetricBasis, c::AbstractVector{<:Number}) = + zeros(eltype(basis.A2Bmap), size(basis.A2Bmap, 2)) _alloc_dAco(dAAdA::AbstractVector, A::AbstractVector, c̃::AbstractArray, args...) = - _alloc_dAco(dAAdA, A, c̃[1], args...) - -function _alloc_dAco(dAAdA::AbstractVector, A::AbstractVector, - c̃::Union{TP, SVector{N, TP}}, dp = _One() - ) where {N, TP <: AbstractProperty} - c̃_dp = contract(c̃, dp) - @show eltype(dAAdA) - @show typeof(c̃_dp) - @show promote_type(eltype(dAAdA), eltype(c̃_dp)) - zeros( promote_type(eltype(dAAdA), typeof(c̃_dp)), length(A) ) + _alloc_dAco(dAAdA, A, c̃[1], args...) + +function _alloc_dAco(dAAdA::AbstractVector, A::AbstractVector, + c̃::Union{TP,SVector{N,TP}}, dp=_One() +) where {N,TP<:AbstractProperty} + c̃_dp = contract(c̃, dp) + @show eltype(dAAdA) + @show typeof(c̃_dp) + @show promote_type(eltype(dAAdA), eltype(c̃_dp)) + zeros(promote_type(eltype(dAAdA), typeof(c̃_dp)), length(A)) end # ------------------------------------------------------------ @@ -114,36 +114,36 @@ end -evaluate(::LinearACEModel, V::ProductEvaluator, cfg::AbstractConfiguration) = - evaluate(V::ProductEvaluator, cfg) +evaluate(::LinearACEModel, V::ProductEvaluator, cfg::AbstractConfiguration) = + evaluate(V::ProductEvaluator, cfg) # compute one "site energy" function evaluate(V::ProductEvaluator, cfg::AbstractConfiguration) - A = evaluate(V.pibasis.basis1p, cfg) - spec = V.pibasis.spec - pireal = V.pibasis.real - symreal = V.real - # initialize output with a sensible type - val = symreal(zero(eltype(V.coeffs)) * pireal(zero(eltype(A)))) - - # constant (0-order) - if spec.orders[1] == 0 - val += V.coeffs[1] - iAAinit = 2 - else - iAAinit = 1 - end - - @inbounds for iAA = iAAinit:length(spec) - aa = A[spec.iAA2iA[iAA, 1]] - for t = 2:spec.orders[iAA] - aa *= A[spec.iAA2iA[iAA, t]] - end - val += symreal(pireal(aa) * V.coeffs[iAA]) - end - - release!(A) - return val + A = evaluate(V.pibasis.basis1p, cfg) + spec = V.pibasis.spec + pireal = V.pibasis.real + symreal = V.real + # initialize output with a sensible type + val = symreal(zero(eltype(V.coeffs)) * pireal(zero(eltype(A)))) + + # constant (0-order) + if spec.orders[1] == 0 + val += V.coeffs[1] + iAAinit = 2 + else + iAAinit = 1 + end + + @inbounds for iAA = iAAinit:length(spec) + aa = A[spec.iAA2iA[iAA, 1]] + for t = 2:spec.orders[iAA] + aa *= A[spec.iAA2iA[iAA, t]] + end + val += symreal(pireal(aa) * V.coeffs[iAA]) + end + + release!(A) + return val end @@ -151,18 +151,18 @@ end -# #for multiple properties. dispatch on the pullback input being a matrix. -# #Basically the same code, except for some parts where we loop over all properties. -# #We generate a list of size "nprop" and keep the same objects as for a single property +# #for multiple properties. dispatch on the pullback input being a matrix. +# #Basically the same code, except for some parts where we loop over all properties. +# #We generate a list of size "nprop" and keep the same objects as for a single property # #inside the list. # function adjoint_EVAL_D(m::LinearACEModel, V::ProductEvaluator, cfg, wt::Matrix) -# _contract = ACEfrictionCore.contract - +# _contract = ACEfrictionCore.contract + # basis1p = V.pibasis.basis1p -# dAAdA = zero(MVector{10, ComplexF64}) # TODO: VERY RISKY -> FIX THIS +# dAAdA = zero(MVector{10, ComplexF64}) # TODO: VERY RISKY -> FIX THIS # A = zeros(ComplexF64, length(basis1p)) # TDX = gradtype(m.basis, cfg) -# dA = zeros(complex(TDX) , length(A), length(cfg)) +# dA = zeros(complex(TDX) , length(A), length(cfg)) # _real = V.real # dAAw = [acquire_B!(V.pibasis, cfg) for _ in 1:length(m.c[1])] # dAw = [similar(A) for _ in 1:length(m.c[1])] @@ -180,13 +180,13 @@ end # end # end -# # [2] dAA_k +# # [2] dAA_k # spec = V.pibasis.spec # for i in 1:length(m.c[1]) # fill!(dAAw[i], 0) # end # for prop in 1:length(m.c[1]) -# if spec.orders[1] == 0; iAAinit=2; else; iAAinit=1; end +# if spec.orders[1] == 0; iAAinit=2; else; iAAinit=1; end # @inbounds for iAA = iAAinit:length(spec) # _AA_local_adjoints!(dAAdA, A, spec.iAA2iA, iAA, spec.orders[iAA], _real) # @fastmath for t = 1:spec.orders[iAA] @@ -204,4 +204,4 @@ end # # [3] dB_k # return dB -# end \ No newline at end of file +# end diff --git a/src/linearmodel.jl b/src/linearmodel.jl index 98be86b..a489878 100644 --- a/src/linearmodel.jl +++ b/src/linearmodel.jl @@ -1,111 +1,111 @@ # -# Draft ACEModel Interface: +# Draft ACEModel Interface: # -# * A model incorporates both the parameterisation and the parameters. +# * A model incorporates both the parameterisation and the parameters. # * E.g. a LinearACEModel would know about the basis and the coefficients -# * We can then perform the following operations: +# * We can then perform the following operations: # - evaluate the model at a given configuration, with given parameters -# - take gradient w.r.t. the configuration -# - take gradient w.r.t. the parameters +# - take gradient w.r.t. the configuration +# - take gradient w.r.t. the parameters # -# An advantage will be that we can do the parameter reduction trick and then -# use a fast evaluator to obtain the gradient w.r.t. configuration which is -# very expensive if we naively take derivatives w.r.t. the basis and then +# An advantage will be that we can do the parameter reduction trick and then +# use a fast evaluator to obtain the gradient w.r.t. configuration which is +# very expensive if we naively take derivatives w.r.t. the basis and then # apply parameters. (this is not implemented below and will be the next step) -# Another advantage is that model, and full model construction are all stored +# Another advantage is that model, and full model construction are all stored # together for future inspection. """ -`struct LinearACEModel`: linear model for symmetric properties in terms of -a `SymmetricBasis`. +`struct LinearACEModel`: linear model for symmetric properties in terms of +a `SymmetricBasis`. -The typical way to construct a linear model is to first construct a basis -`basis`, some default coefficients `c` and then call +The typical way to construct a linear model is to first construct a basis +`basis`, some default coefficients `c` and then call ```julia model = LinearACEModel(basis, c) ``` -### Multiple properties +### Multiple properties -If `c::Vector{<: Number}` then the output of the model will be the property -encoded in the basis. But one can also use a single basis to produce -multiple properties (with different coefficients). This can be achieved by -simply supplying `c::Vector{SVector{N, T}}` where `N` will then be the -number of properties. +If `c::Vector{<: Number}` then the output of the model will be the property +encoded in the basis. But one can also use a single basis to produce +multiple properties (with different coefficients). This can be achieved by +simply supplying `c::Vector{SVector{N, T}}` where `N` will then be the +number of properties. """ -struct LinearACEModel{TB, TP, TEV} <: AbstractACEModel - basis::TB - c::Vector{TP} - evaluator::TEV -end - -struct NaiveEvaluator end - -function LinearACEModel(basis::SymmetricBasis, c = zeros(length(basis)); - evaluator = :standard) - if evaluator == :naive - ev = NaiveEvaluator() - elseif evaluator == :standard - ev = ProductEvaluator(basis, c) - elseif evaluator == :recursive - error("Recursive evaluator not yet implemented") - else - error("unknown evaluator") - end - return LinearACEModel(basis, c, ev) +struct LinearACEModel{TB,TP,TEV} <: AbstractACEModel + basis::TB + c::Vector{TP} + evaluator::TEV end -# LinearACEModel(basis::SymmetricBasis, c::Vector, evaluator) = +struct NaiveEvaluator end + +function LinearACEModel(basis::SymmetricBasis, c=zeros(length(basis)); + evaluator=:standard) + if evaluator == :naive + ev = NaiveEvaluator() + elseif evaluator == :standard + ev = ProductEvaluator(basis, c) + elseif evaluator == :recursive + error("Recursive evaluator not yet implemented") + else + error("unknown evaluator") + end + return LinearACEModel(basis, c, ev) +end + +# LinearACEModel(basis::SymmetricBasis, c::Vector, evaluator) = # LinearACEModel(basis, c, ev, VectorPool{eltype(c)}) -# ------- parameter wrangling +# ------- parameter wrangling nparams(m::LinearACEModel) = length(m.c) params(m::LinearACEModel) = copy(m.c) -function set_params!(m::LinearACEModel, c) - m.c[:] .= c - set_params!(m.evaluator, m.basis, c) - return m +function set_params!(m::LinearACEModel, c) + m.c[:] .= c + set_params!(m.evaluator, m.basis, c) + return m end -set_params!(::NaiveEvaluator, args...) = nothing +set_params!(::NaiveEvaluator, args...) = nothing # ------------------- FIO -==(V1::LinearACEModel, V2::LinearACEModel) = - _allfieldsequal(V1, V2) - -write_dict(V::LinearACEModel) = - Dict( "__id__" => "ACEfrictionCore_LinearACEModel", - "basis" => write_dict(V.basis), - "c" => write_dict(V.c), - "evaluator" => write_dict(V.evaluator) ) - -function read_dict(::Val{:ACEfrictionCore_LinearACEModel}, D::Dict) - basis = read_dict(D["basis"]) - c = read_dict(D["c"]) - # special evaluator version of the read_dict - evaluator = read_dict(Val(Symbol(D["evaluator"]["__id__"])), - D["evaluator"], basis, c) - return LinearACEModel(basis, c, evaluator) +==(V1::LinearACEModel, V2::LinearACEModel) = + _allfieldsequal(V1, V2) + +write_dict(V::LinearACEModel) = + Dict("__id__" => "ACEfrictionCore_LinearACEModel", + "basis" => write_dict(V.basis), + "c" => write_dict(V.c), + "evaluator" => write_dict(V.evaluator)) + +function read_dict(::Val{:ACEfrictionCore_LinearACEModel}, D::AbstractDict) + basis = read_dict(D["basis"]) + c = read_dict(D["c"]) + # special evaluator version of the read_dict + evaluator = read_dict(Val(Symbol(D["evaluator"]["__id__"])), + D["evaluator"], basis, c) + return LinearACEModel(basis, c, evaluator) end -write_dict(ev::NaiveEvaluator) = - Dict("__id__" => "ACEfrictionCore_NaiveEvaluator" ) +write_dict(ev::NaiveEvaluator) = + Dict("__id__" => "ACEfrictionCore_NaiveEvaluator") -read_dict(::Val{:ACEfrictionCore_NaiveEvaluator}, D::Dict, args...) = - NaiveEvaluator() +read_dict(::Val{:ACEfrictionCore_NaiveEvaluator}, D::AbstractDict, args...) = + NaiveEvaluator() -# ------------------- dispatching on the evaluators +# ------------------- dispatching on the evaluators -evaluate(m::LinearACEModel, X::AbstractConfiguration) = - evaluate(m, m.evaluator, X) +evaluate(m::LinearACEModel, X::AbstractConfiguration) = + evaluate(m, m.evaluator, X) # grad_config function removed - derivative functionality has been removed @@ -115,34 +115,34 @@ evaluate(m::LinearACEModel, X::AbstractConfiguration) = - # TODO: fix terminology, bring in linear with the _rrule_.... thing -adjoint_EVAL_D(m::LinearACEModel, cfg::AbstractConfiguration, w) = - adjoint_EVAL_D(m, m.evaluator, cfg, w) +# TODO: fix terminology, bring in linear with the _rrule_.... thing +adjoint_EVAL_D(m::LinearACEModel, cfg::AbstractConfiguration, w) = + adjoint_EVAL_D(m, m.evaluator, cfg, w) -# ------------------- implementation of naive evaluator -# this is only intended for testing, as it uses the naive evaluation of +# ------------------- implementation of naive evaluator +# this is only intended for testing, as it uses the naive evaluation of # the symmetric basis, rather than the conversion to the AA basis -function evaluate(m::LinearACEModel, ::NaiveEvaluator, cfg::AbstractConfiguration) - B = evaluate(m.basis, cfg) - val = contract(m.c, B) - release!(B) - return val -end +function evaluate(m::LinearACEModel, ::NaiveEvaluator, cfg::AbstractConfiguration) + B = evaluate(m.basis, cfg) + val = contract(m.c, B) + release!(B) + return val +end # grad_config function for NaiveEvaluator removed - derivative functionality has been removed -function adjoint_EVAL_D(m::LinearACEModel, ::NaiveEvaluator, - X::AbstractConfiguration, w) - error("grad_params_config functionality has been removed") - # dB = grad_params_config(m, X) -- removed - # g = zeros(size(dB, 1)) -- removed - # for i = 1:length(g), j = 1:size(dB, 2) -- removed - # g[i] += dot(dB[i, j], w[j]) -- removed - # end -- removed - # release_dB!(m.basis, dB) -- removed - # return g -- removed +function adjoint_EVAL_D(m::LinearACEModel, ::NaiveEvaluator, + X::AbstractConfiguration, w) + error("grad_params_config functionality has been removed") + # dB = grad_params_config(m, X) -- removed + # g = zeros(size(dB, 1)) -- removed + # for i = 1:length(g), j = 1:size(dB, 2) -- removed + # g[i] += dot(dB[i, j], w[j]) -- removed + # end -- removed + # release_dB!(m.basis, dB) -- removed + # return g -- removed end @@ -152,10 +152,10 @@ end function _adj_evaluate(dp, model::ACEfrictionCore.LinearACEModel, cfg) - error("grad_params functionality has been removed") - # gp_ = ACEfrictionCore.grad_params(model, cfg) -- removed - # gp = [ a * dp for a in gp_ ] -- removed - # return NoTangent(), gp, _rrule_evaluate(dp, model, cfg) -- removed + error("grad_params functionality has been removed") + # gp_ = ACEfrictionCore.grad_params(model, cfg) -- removed + # gp = [ a * dp for a in gp_ ] -- removed + # return NoTangent(), gp, _rrule_evaluate(dp, model, cfg) -- removed end # ChainRules.rrule function removed - derivative functionality has been removed diff --git a/src/pibasis.jl b/src/pibasis.jl index 2d50161..f243527 100644 --- a/src/pibasis.jl +++ b/src/pibasis.jl @@ -8,8 +8,8 @@ export PIBasis `struct PIBasisSpec` """ struct PIBasisSpec - orders::Vector{Int} # order (length) of ith basis function - iAA2iA::Matrix{Int} # where in A can we find the ith basis function + orders::Vector{Int} # order (length) of ith basis function + iAA2iA::Matrix{Int} # where in A can we find the ith basis function end ==(B1::PIBasisSpec, B2::PIBasisSpec) = _allfieldsequal(B1, B2) @@ -19,111 +19,111 @@ Base.length(spec::PIBasisSpec) = length(spec.orders) maxcorrorder(spec::PIBasisSpec) = size(spec.iAA2iA, 2) function _get_pibfcn(spec0, Aspec, vv) - vv1 = vv[2:end] - vvnz = vv1[findall(vv1 .!= 0)] - return (spec0[vv[1]], Aspec[vvnz]) + vv1 = vv[2:end] + vvnz = vv1[findall(vv1 .!= 0)] + return (spec0[vv[1]], Aspec[vvnz]) end function _get_pibfcn(Aspec, vv) - vvnz = vv[findall(vv .!= 0)] - return Aspec[vvnz] + vvnz = vv[findall(vv .!= 0)] + return Aspec[vvnz] end -# TODO: maybe instead of property == nothing, there should be a -# generic property with no symmetry attached to it. - -function PIBasisSpec( basis1p::OneParticleBasis, - symgrp::SymmetryGroup, - Bsel::DownsetBasisSelector; - property = nothing, - filterfun = _->true, - init1pbasis = true ) - - # we initialize the 1p-basis here; to prevent this it must be manually - # avoided by passing in init1pbasis = false - if init1pbasis - init1pspec!(basis1p, Bsel) - end - - # get the basis spec of the one-particle basis - # Aspec[i] described the basis function that will get written into A[i] - Aspec = get_spec(basis1p) - - # we assume that `Aspec` is sorted by degree, but best to double-check this - # since the notion of degree used to construct `Aspec` might be different - # from the one used to construct AAspec. - if !issorted(Aspec; by = b -> level(b, Bsel, basis1p)) - error("""PIBasisSpec : AAspec construction failed because Aspec is not - sorted by degree. This could e.g. happen if an incompatible - notion of degree was used to construct the 1-p basis spec.""") - end - # An AA basis function is given by a tuple 𝒗 = vv. Each index 𝒗ᵢ = vv[i] - # corresponds to the basis function Aspec[𝒗ᵢ] and the tuple - # 𝒗 = (𝒗₁, ...) to a product basis function - # ∏ A_{vₐ} - tup2b = vv -> _get_pibfcn(Aspec, vv) - - # degree or level of a basis function ↦ is it admissible? - admissible = bb -> (level(bb, Bsel, basis1p) <= maxlevel(bb, Bsel, basis1p)) - - if property != nothing - filter1 = bb -> filterfun(bb) && filter(bb, Bsel, basis1p) && filter(property, symgrp, bb) - else - filter1 = bb -> filterfun(bb) && filter(bb, Bsel, basis1p) - end - - - # we can now construct the basis specification; the `ordered = true` - # keyword signifies that this is a permutation-invariant basis - maxord = maxorder(Bsel) - AAspec = gensparse(; NU = maxorder(Bsel), - tup2b = tup2b, - admissible = admissible, - ordered = true, - maxvv = [length(Aspec) for _=1:maxord], - filter = filter1) - - return PIBasisSpec(AAspec) +# TODO: maybe instead of property == nothing, there should be a +# generic property with no symmetry attached to it. + +function PIBasisSpec(basis1p::OneParticleBasis, + symgrp::SymmetryGroup, + Bsel::DownsetBasisSelector; + property=nothing, + filterfun=_ -> true, + init1pbasis=true) + + # we initialize the 1p-basis here; to prevent this it must be manually + # avoided by passing in init1pbasis = false + if init1pbasis + init1pspec!(basis1p, Bsel) + end + + # get the basis spec of the one-particle basis + # Aspec[i] described the basis function that will get written into A[i] + Aspec = get_spec(basis1p) + + # we assume that `Aspec` is sorted by degree, but best to double-check this + # since the notion of degree used to construct `Aspec` might be different + # from the one used to construct AAspec. + if !issorted(Aspec; by=b -> level(b, Bsel, basis1p)) + error("""PIBasisSpec : AAspec construction failed because Aspec is not + sorted by degree. This could e.g. happen if an incompatible + notion of degree was used to construct the 1-p basis spec.""") + end + # An AA basis function is given by a tuple 𝒗 = vv. Each index 𝒗ᵢ = vv[i] + # corresponds to the basis function Aspec[𝒗ᵢ] and the tuple + # 𝒗 = (𝒗₁, ...) to a product basis function + # ∏ A_{vₐ} + tup2b = vv -> _get_pibfcn(Aspec, vv) + + # degree or level of a basis function ↦ is it admissible? + admissible = bb -> (level(bb, Bsel, basis1p) <= maxlevel(bb, Bsel, basis1p)) + + if property != nothing + filter1 = bb -> filterfun(bb) && filter(bb, Bsel, basis1p) && filter(property, symgrp, bb) + else + filter1 = bb -> filterfun(bb) && filter(bb, Bsel, basis1p) + end + + + # we can now construct the basis specification; the `ordered = true` + # keyword signifies that this is a permutation-invariant basis + maxord = maxorder(Bsel) + AAspec = gensparse(; NU=maxorder(Bsel), + tup2b=tup2b, + admissible=admissible, + ordered=true, + maxvv=[length(Aspec) for _ = 1:maxord], + filter=filter1) + + return PIBasisSpec(AAspec) end function PIBasisSpec(AAspec) - orders = zeros(Int, length(AAspec)) - iAA2iA = zeros(Int, (length(AAspec), length(AAspec[1]))) - for (iAA, vv) in enumerate(AAspec) - # we use reverse because gensparse constructs the indices in - # ascending order, but we want descending here. - # (I don't remember why though) - iAA2iA[iAA, :] .= reverse(vv) - orders[iAA] = length( findall( vv .!= 0 ) ) - end - return PIBasisSpec(orders, iAA2iA) + orders = zeros(Int, length(AAspec)) + iAA2iA = zeros(Int, (length(AAspec), length(AAspec[1]))) + for (iAA, vv) in enumerate(AAspec) + # we use reverse because gensparse constructs the indices in + # ascending order, but we want descending here. + # (I don't remember why though) + iAA2iA[iAA, :] .= reverse(vv) + orders[iAA] = length(findall(vv .!= 0)) + end + return PIBasisSpec(orders, iAA2iA) end get_spec(AAspec::PIBasisSpec, i::Integer) = AAspec.iAA2iA[i, 1:AAspec.orders[i]] -# ------------------ PISpec sparsification +# ------------------ PISpec sparsification """ -returns a new `PIBasisSpec` constructed from the old one, but keeping only -the basis indices `Ikeep`. This maintains the order of the basis functions. +returns a new `PIBasisSpec` constructed from the old one, but keeping only +the basis indices `Ikeep`. This maintains the order of the basis functions. """ -sparsify(spec::PIBasisSpec, Ikeep::AbstractVector{<: Integer}) = - PIBasisSpec(spec.orders[Ikeep], spec.iAA2iA[Ikeep, :]) +sparsify(spec::PIBasisSpec, Ikeep::AbstractVector{<:Integer}) = + PIBasisSpec(spec.orders[Ikeep], spec.iAA2iA[Ikeep, :]) -function _fix_A_indices!(spec::PIBasisSpec, new_inds::AbstractVector{<: Integer}) - for iAA = 1:size(spec.iAA2iA, 1) - for α = 1:spec.orders[iAA] - vα = spec.iAA2iA[iAA, α] - new_vα = new_inds[vα] - @assert new_vα > 0 - spec.iAA2iA[iAA, α] = new_vα - end - end - return nothing +function _fix_A_indices!(spec::PIBasisSpec, new_inds::AbstractVector{<:Integer}) + for iAA = 1:size(spec.iAA2iA, 1) + for α = 1:spec.orders[iAA] + vα = spec.iAA2iA[iAA, α] + new_vα = new_inds[vα] + @assert new_vα > 0 + spec.iAA2iA[iAA, α] = new_vα + end + end + return nothing end # --------------------------------- PIBasis implementation @@ -142,68 +142,68 @@ PIBasis(basis1p, N, D, maxdeg) * `D` : an abstract degee specification, e.g., SparsePSHDegree * `maxdeg` : the maximum polynomial degree as measured by `D` """ -mutable struct PIBasis{BOP, REAL} <: ACEBasis - basis1p::BOP # a one-particle basis - spec::PIBasisSpec - real::REAL # could be `real` or `identity` to keep AA complex -end # TODO -> chain?? +mutable struct PIBasis{BOP,REAL} <: ACEBasis + basis1p::BOP # a one-particle basis + spec::PIBasisSpec + real::REAL # could be `real` or `identity` to keep AA complex +end # TODO -> chain?? cutoff(basis::PIBasis) = cutoff(basis.basis1p) -==(B1::PIBasis, B2::PIBasis) = - ( (B1.basis1p == B2.basis1p) && - (B1.spec == B2.spec) && - (B1.real == B2.real) ) +==(B1::PIBasis, B2::PIBasis) = + ((B1.basis1p == B2.basis1p) && + (B1.spec == B2.spec) && + (B1.real == B2.real)) Base.length(basis::PIBasis) = length(basis.spec) -# default symmetry group -PIBasis(basis1p, Bsel::AbstractBasisSelector; kwargs...) = - PIBasis(basis1p, O3(), Bsel; kwargs...) +# default symmetry group +PIBasis(basis1p, Bsel::AbstractBasisSelector; kwargs...) = + PIBasis(basis1p, O3(), Bsel; kwargs...) -PIBasis(basis1p, symgrp, Bsel::AbstractBasisSelector; - isreal = false, kwargs...) = - PIBasis(basis1p, - PIBasisSpec(basis1p, symgrp, Bsel; kwargs...), - isreal ? Base.real : Base.identity ) +PIBasis(basis1p, symgrp, Bsel::AbstractBasisSelector; + isreal=false, kwargs...) = + PIBasis(basis1p, + PIBasisSpec(basis1p, symgrp, Bsel; kwargs...), + isreal ? Base.real : Base.identity) get_spec(pibasis::PIBasis) = - [ get_spec(pibasis, i) for i = 1:length(pibasis) ] + [get_spec(pibasis, i) for i = 1:length(pibasis)] get_spec(pibasis::PIBasis, i::Integer) = - get_spec.( Ref(pibasis.basis1p), get_spec(pibasis.spec, i) ) + get_spec.(Ref(pibasis.basis1p), get_spec(pibasis.spec, i)) setreal(basis::PIBasis, isreal::Bool) = - PIBasis(basis.basis1p, basis.spec, isreal) + PIBasis(basis.basis1p, basis.spec, isreal) maxcorrorder(basis::PIBasis) = maxcorrorder(basis.spec) -# ------------------ sparsification +# ------------------ sparsification -function sparsify!(basis::PIBasis, Ikeep::AbstractVector{<: Integer}) - basis.spec = sparsify(basis.spec, Ikeep) - return basis -end +function sparsify!(basis::PIBasis, Ikeep::AbstractVector{<:Integer}) + basis.spec = sparsify(basis.spec, Ikeep) + return basis +end """ -This should allow the 1p basis to sparsify itself, then feed back to the +This should allow the 1p basis to sparsify itself, then feed back to the pibasis what the correct indices are. """ function clean_1pbasis!(basis::PIBasis) - spec = get_spec(basis) - B1p = basis.basis1p - spec1p = NamedTuple[] - for bb in spec - append!(spec1p, bb) - end - identity.(unique!(spec1p)) - # sparsify the product 1p basis - _, new_inds = sparsify!(basis.basis1p, spec1p) - # now fix the indexing of the PIBasis specification - _fix_A_indices!(basis.spec, new_inds) - return basis + spec = get_spec(basis) + B1p = basis.basis1p + spec1p = NamedTuple[] + for bb in spec + append!(spec1p, bb) + end + identity.(unique!(spec1p)) + # sparsify the product 1p basis + _, new_inds = sparsify!(basis.basis1p, spec1p) + # now fix the indexing of the PIBasis specification + _fix_A_indices!(basis.spec, new_inds) + return basis end @@ -216,15 +216,15 @@ _scaling_absvalue(x::Symbol) = 0 # TODO: this is a hack; cf. #68 function scaling(pibasis::PIBasis, p) - ww = zeros(Float64, length(pibasis)) - bspec = get_spec(pibasis) - for i = 1:length(pibasis) - for b in bspec[i] - # TODO: revisit how this should be implemented for a general basis - ww[i] += sum(x -> _scaling_absvalue(x)^p, b) # abs.(values(b)).^p - end - end - return ww + ww = zeros(Float64, length(pibasis)) + bspec = get_spec(pibasis) + for i = 1:length(pibasis) + for b in bspec[i] + # TODO: revisit how this should be implemented for a general basis + ww[i] += sum(x -> _scaling_absvalue(x)^p, b) # abs.(values(b)).^p + end + end + return ww end @@ -233,123 +233,120 @@ end # FIO codes write_dict(basis::PIBasis) = - Dict( "__id__" => "ACEfrictionCore_PIBasis", - "basis1p" => write_dict(basis.basis1p), - "spec" => write_dict(basis.spec), - "real" => basis.real == Base.real ? true : false ) + Dict("__id__" => "ACEfrictionCore_PIBasis", + "basis1p" => write_dict(basis.basis1p), + "spec" => write_dict(basis.spec), + "real" => basis.real == Base.real ? true : false) -read_dict(::Val{:ACEfrictionCore_PIBasis}, D::Dict) = - PIBasis( read_dict(D["basis1p"]), - read_dict(D["spec"]), - D["real"] ? Base.real : Base.identity ) +read_dict(::Val{:ACEfrictionCore_PIBasis}, D::AbstractDict) = + PIBasis(read_dict(D["basis1p"]), + read_dict(D["spec"]), + D["real"] ? Base.real : Base.identity) write_dict(spec::PIBasisSpec) = - Dict( "__id__" => "ACEfrictionCore_PIBasisSpec", - "orders" => spec.orders, - "iAA2iA" => write_dict(spec.iAA2iA) ) + Dict("__id__" => "ACEfrictionCore_PIBasisSpec", + "orders" => spec.orders, + "iAA2iA" => write_dict(spec.iAA2iA)) -read_dict(::Val{:ACEfrictionCore_PIBasisSpec}, D::Dict) = - PIBasisSpec( D["orders"], read_dict(D["iAA2iA"]) ) +read_dict(::Val{:ACEfrictionCore_PIBasisSpec}, D::AbstractDict) = + PIBasisSpec(D["orders"], read_dict(D["iAA2iA"])) # ------------------------------------------------- # Evaluation codes function evaluate!(AA, basis::PIBasis, config::UConfig) - A = evaluate(basis.basis1p, config) - evaluate!(AA, basis, A) - release!(A) - return AA + A = evaluate(basis.basis1p, config) + evaluate!(AA, basis, A) + release!(A) + return AA end -function evaluate!(AA, basis::PIBasis, A::AbstractVector{<: Number}) - fill!(AA, 1) - for iAA = 1:length(basis) - aa = one(eltype(A)) - for t = 1:basis.spec.orders[iAA] - aa *= A[ basis.spec.iAA2iA[ iAA, t ] ] - end - AA[iAA] = basis.real(aa) - end - return AA +function evaluate!(AA, basis::PIBasis, A::AbstractVector{<:Number}) + fill!(AA, 1) + for iAA = 1:length(basis) + aa = one(eltype(A)) + for t = 1:basis.spec.orders[iAA] + aa *= A[basis.spec.iAA2iA[iAA, t]] + end + AA[iAA] = basis.real(aa) + end + return AA end -_valtype(basis::PIBasis, A::AbstractVector{<: Number}) = - basis.real(eltype(A)) +_valtype(basis::PIBasis, A::AbstractVector{<:Number}) = + basis.real(eltype(A)) # draft of defining bases via chains -function evaluate(basis::PIBasis, A::AbstractVector{<: Number}) - VT = _valtype(basis, A) - AA = Vector{VT}(undef, length(basis)) - evaluate!(AA, basis, A) - return AA +function evaluate(basis::PIBasis, A::AbstractVector{<:Number}) + VT = _valtype(basis, A) + AA = Vector{VT}(undef, length(basis)) + evaluate!(AA, basis, A) + return AA end # draft of defining bases via chains -function evaluate(basis::PIBasis, config::UConfig) - A = evaluate(basis.basis1p, config) - AA = evaluate(basis, A) - release!(A) - return AA +function evaluate(basis::PIBasis, config::UConfig) + A = evaluate(basis.basis1p, config) + AA = evaluate(basis, A) + release!(A) + return AA end function _AA_local_adjoints!(dAAdA, A, iAA2iA, iAA, ord, _real) - if ord == 1 - return _AA_local_adjoints_1!(dAAdA, A, iAA2iA, iAA, ord, _real) - elseif ord == 2 - return _AA_local_adjoints_2!(dAAdA, A, iAA2iA, iAA, ord, _real) - else - return _AA_local_adjoints_x!(dAAdA, A, iAA2iA, iAA, ord, _real) - end + if ord == 1 + return _AA_local_adjoints_1!(dAAdA, A, iAA2iA, iAA, ord, _real) + elseif ord == 2 + return _AA_local_adjoints_2!(dAAdA, A, iAA2iA, iAA, ord, _real) + else + return _AA_local_adjoints_x!(dAAdA, A, iAA2iA, iAA, ord, _real) + end end function _AA_local_adjoints_1!(dAAdA, A, iAA2iA, iAA, ord, _real) - @inbounds dAAdA[1] = 1 - @inbounds A1 = A[iAA2iA[iAA, 1]] - return _real(A1) + @inbounds dAAdA[1] = 1 + @inbounds A1 = A[iAA2iA[iAA, 1]] + return _real(A1) end function _AA_local_adjoints_2!(dAAdA, A, iAA2iA, iAA, ord, _real) - @inbounds A1 = A[iAA2iA[iAA, 1]] - @inbounds A2 = A[iAA2iA[iAA, 2]] - @inbounds dAAdA[1] = A2 - @inbounds dAAdA[2] = A1 - return _real(A1 * A2) + @inbounds A1 = A[iAA2iA[iAA, 1]] + @inbounds A2 = A[iAA2iA[iAA, 2]] + @inbounds dAAdA[1] = A2 + @inbounds dAAdA[2] = A1 + return _real(A1 * A2) end function _AA_local_adjoints_x!(dAAdA, A, iAA2iA, iAA, ord, _real) - @assert length(dAAdA) >= ord - @assert ord >= 2 - # TODO - optimize a bit more? can move one operation out of the loop - # Forward pass: - @inbounds A1 = A[iAA2iA[iAA, 1]] - @inbounds A2 = A[iAA2iA[iAA, 2]] - @inbounds dAAdA[1] = 1 - @inbounds dAAdA[2] = A1 - @inbounds AAfwd = A1 * A2 - @inbounds for a = 3:ord-1 - dAAdA[a] = AAfwd - AAfwd *= A[iAA2iA[iAA, a]] - end - @inbounds dAAdA[ord] = AAfwd - @inbounds Aend = A[iAA2iA[iAA, ord]] - aa = _real(AAfwd * Aend) - # backward pass - @inbounds AAbwd = Aend - @inbounds for a = ord-1:-1:3 - dAAdA[a] *= AAbwd - AAbwd *= A[iAA2iA[iAA, a]] - end - dAAdA[2] *= AAbwd - AAbwd *= A2 - dAAdA[1] *= AAbwd - - return aa + @assert length(dAAdA) >= ord + @assert ord >= 2 + # TODO - optimize a bit more? can move one operation out of the loop + # Forward pass: + @inbounds A1 = A[iAA2iA[iAA, 1]] + @inbounds A2 = A[iAA2iA[iAA, 2]] + @inbounds dAAdA[1] = 1 + @inbounds dAAdA[2] = A1 + @inbounds AAfwd = A1 * A2 + @inbounds for a = 3:ord-1 + dAAdA[a] = AAfwd + AAfwd *= A[iAA2iA[iAA, a]] + end + @inbounds dAAdA[ord] = AAfwd + @inbounds Aend = A[iAA2iA[iAA, ord]] + aa = _real(AAfwd * Aend) + # backward pass + @inbounds AAbwd = Aend + @inbounds for a = ord-1:-1:3 + dAAdA[a] *= AAbwd + AAbwd *= A[iAA2iA[iAA, a]] + end + dAAdA[2] *= AAbwd + AAbwd *= A2 + dAAdA[1] *= AAbwd + + return aa end _acquire_dAAdA!(basis::PIBasis, A) = Vector{eltype(A)}(undef, maxcorrorder(basis)) - - - diff --git a/src/polynomials/orthpolys.jl b/src/polynomials/orthpolys.jl index 87234a7..b5a16fc 100644 --- a/src/polynomials/orthpolys.jl +++ b/src/polynomials/orthpolys.jl @@ -7,15 +7,15 @@ using LinearAlgebra: dot import ACEfrictionCore -import ACEfrictionCore: evaluate!, - evaluate, - # frule_evaluate, - read_dict, write_dict, - inv_transform, - ACEBasis, ScalarACEBasis, - acquire!, release!, - ArrayCache, - chain +import ACEfrictionCore: evaluate!, + evaluate, + # frule_evaluate, + read_dict, write_dict, + inv_transform, + ACEBasis, ScalarACEBasis, + acquire!, release!, + ArrayCache, + chain # using ForwardDiff: derivative @@ -26,36 +26,36 @@ import Base: == export orthpolys, transformed_jacobi, discrete_jacobi -# these inner functions have been timed to run at +# these inner functions have been timed to run at # 6.7ns, 9.2ns, 11.7ns => no need to hand-optimise _fcut_inner(pl, tl, pr, tr, t) = (t - tl)^pl * (t - tr)^pr -# _fcut_d_inner(pl, tl, pr, tr, t) = +# _fcut_d_inner(pl, tl, pr, tr, t) = # derivative( t -> _fcut_inner(pl, tl, pr, tr, t), t ) -# _fcut_dd_inner(pl, tl, pr, tr, t) = +# _fcut_dd_inner(pl, tl, pr, tr, t) = # derivative( t -> _fcut_d_inner(pl, tl, pr, tr, t), t ) function _fcut_(pl, tl, pr, tr, t) - if (pl > 0 && t < tl) || (pr > 0 && t > tr) - return zero(t) - end - return _fcut_inner(pl, tl, pr, tr, t) + if (pl > 0 && t < tl) || (pr > 0 && t > tr) + return zero(t) + end + return _fcut_inner(pl, tl, pr, tr, t) end function _fcut_d_(pl, tl, pr, tr, t) - if (pl > 0 && t < tl) || (pr > 0 && t > tr) - return zero(t) - end - return _fcut_d_inner(pl, tl, pr, tr, t) + if (pl > 0 && t < tl) || (pr > 0 && t > tr) + return zero(t) + end + return _fcut_d_inner(pl, tl, pr, tr, t) end function _fcut_dd_(pl, tl, pr, tr, t) - if (pl > 0 && t < tl) || (pr > 0 && t > tr) - return zero(t) - end - return _fcut_dd_inner(pl, tl, pr, tr, t) + if (pl > 0 && t < tl) || (pr > 0 && t > tr) + return zero(t) + end + return _fcut_dd_inner(pl, tl, pr, tr, t) end @@ -75,152 +75,152 @@ an "envelope". This results in the recursion Orthogonality is achieved with respect to a user-specified distribution, which can be either continuous or discrete. -TODO: say more on the distribution! Maybe generalize to non-diagonal +TODO: say more on the distribution! Maybe generalize to non-diagonal inner products? """ struct OrthPolyBasis{T} <: ScalarACEBasis - # ----------------- the parameters for the cutoff function - pl::Int # cutoff power left - tl::T # cutoff left (transformed variable) - pr::Int # cutoff power right - tr::T # cutoff right (transformed variable) - # ----------------- the recursion coefficients - A::Vector{T} - B::Vector{T} - C::Vector{T} - # ----------------- used only for construction ... - # but useful to have since it defines the notion of orth. - tdf::Vector{T} - ww::Vector{T} - # ------------- - B_pool::ArrayCache{T} - dB_pool::ArrayCache{T} + # ----------------- the parameters for the cutoff function + pl::Int # cutoff power left + tl::T # cutoff left (transformed variable) + pr::Int # cutoff power right + tr::T # cutoff right (transformed variable) + # ----------------- the recursion coefficients + A::Vector{T} + B::Vector{T} + C::Vector{T} + # ----------------- used only for construction ... + # but useful to have since it defines the notion of orth. + tdf::Vector{T} + ww::Vector{T} + # ------------- + B_pool::ArrayCache{T} + dB_pool::ArrayCache{T} end -OrthPolyBasis(pl, tl::T, pr, tr::T, A::Vector{T}, B::Vector{T}, C::Vector{T}, - tdf, ww) where {T} = - OrthPolyBasis(pl, tl, pr, tr, A, B, C, tdf, ww, - ArrayCache{T}(), ArrayCache{T}()) +OrthPolyBasis(pl, tl::T, pr, tr::T, A::Vector{T}, B::Vector{T}, C::Vector{T}, + tdf, ww) where {T} = + OrthPolyBasis(pl, tl, pr, tr, A, B, C, tdf, ww, + ArrayCache{T}(), ArrayCache{T}()) Base.length(P::OrthPolyBasis) = length(P.A) -_valtype(::OrthPolyBasis{T}, t::S) where {T, S} = - promote_type(T, S) +_valtype(::OrthPolyBasis{T}, t::S) where {T,S} = + promote_type(T, S) ==(J1::OrthPolyBasis, J2::OrthPolyBasis) = - all( getfield(J1, sym) == getfield(J2, sym) - for sym in (:pr, :tr, :pl, :tl, :A, :B, :C) ) + all(getfield(J1, sym) == getfield(J2, sym) + for sym in (:pr, :tr, :pl, :tl, :A, :B, :C)) -Base.show(io::IO, P::OrthPolyBasis) = - print(io, "OrthPolyBasis(pl = $(P.pl), tl = $(P.tl), pr = $(P.pr), tr = $(P.tr), ...)") +Base.show(io::IO, P::OrthPolyBasis) = + print(io, "OrthPolyBasis(pl = $(P.pl), tl = $(P.tl), pr = $(P.pr), tr = $(P.tr), ...)") write_dict(J::OrthPolyBasis{T}) where {T} = Dict( - "__id__" => "ACEfrictionCore_OrthPolyBasis", - "T" => write_dict(T), - "pr" => J.pr, - "tr" => J.tr, - "pl" => J.pl, - "tl" => J.tl, - "A" => J.A, - "B" => J.B, - "C" => J.C, - "tdf" => J.tdf, - "ww" => J.ww - ) - -OrthPolyBasis(D::Dict, T=read_dict(D["T"])) = - OrthPolyBasis( - D["pl"], D["tl"], D["pr"], D["tr"], - Vector{T}(D["A"]), Vector{T}(D["B"]), Vector{T}(D["C"]), - T.(D["tdf"]), T.(D["ww"]) - ) - -read_dict(::Val{:ACEfrictionCore_OrthPolyBasis}, D::Dict) = OrthPolyBasis(D) + "__id__" => "ACEfrictionCore_OrthPolyBasis", + "T" => write_dict(T), + "pr" => J.pr, + "tr" => J.tr, + "pl" => J.pl, + "tl" => J.tl, + "A" => J.A, + "B" => J.B, + "C" => J.C, + "tdf" => J.tdf, + "ww" => J.ww +) + +OrthPolyBasis(D::AbstractDict, T=read_dict(D["T"])) = + OrthPolyBasis( + D["pl"], D["tl"], D["pr"], D["tr"], + Vector{T}(D["A"]), Vector{T}(D["B"]), Vector{T}(D["C"]), + T.(D["tdf"]), T.(D["ww"]) + ) + +read_dict(::Val{:ACEfrictionCore_OrthPolyBasis}, D::AbstractDict) = OrthPolyBasis(D) # rand applied to a J will return a random transformed distance drawn from # the measure w.r.t. which the polynomials were constructed. # TODO: allow non-constant weights! function ACEfrictionCore.rand_radial(J::OrthPolyBasis) - @assert maximum(abs, diff(J.ww)) == 0 - return rand(J.tdf) + @assert maximum(abs, diff(J.ww)) == 0 + return rand(J.tdf) end -OrthPolyBasis(N::Integer, J::OrthPolyBasis) = - OrthPolyBasis(N, J.pcut, J.tcut, J.pin, J.tin, J.tdf, J.ww) +OrthPolyBasis(N::Integer, J::OrthPolyBasis) = + OrthPolyBasis(N, J.pcut, J.tcut, J.pin, J.tin, J.tdf, J.ww) function OrthPolyBasis(N::Integer, - pcut::Integer, - tcut::T, - pin::Integer, - tin::T, - tdf::AbstractVector{T}, - ww::AbstractVector{T} = ones(T, length(tdf)) - ) where {T <: AbstractFloat} - @assert pcut >= 0 && pin >= 0 - @assert N > 0 - - if tcut < tin - tl, tr = tcut, tin - pl, pr = pcut, pin - else - tl, tr = tin, tcut - pl, pr = pin, pcut - end - - if minimum(tdf) < tl || maximum(tdf) > tr - @warn("OrthPolyBasis: t range outside [tl, tr]") - end - - A = zeros(T, N) - B = zeros(T, N) - C = zeros(T, N) - - # normalise the weights s.t. <1, 1> = 1 - ww = ww ./ sum(ww) - # define inner products - dotw = (f1, f2) -> dot(f1, ww .* f2) - - # start the iteration - _J1 = _fcut_.(pl, tl, pr, tr, tdf) - a = sqrt( dotw(_J1, _J1) ) - A[1] = 1/a - J1 = A[1] * _J1 - - if N > 1 - # a J2 = (t - b) J1 - b = dotw(tdf .* J1, J1) - _J2 = (tdf .- b) .* J1 - a = sqrt( dotw(_J2, _J2) ) - A[2] = 1/a - B[2] = -b / a - J2 = (A[2] * tdf .+ B[2]) .* J1 - - # keep the last two for the 3-term recursion - Jprev = J2 - Jpprev = J1 - end - - for n = 3:N - # a Jn = (t - b) J_{n-1} - c J_{n-2} - b = dotw(tdf .* Jprev, Jprev) - c = dotw(tdf .* Jprev, Jpprev) - _J = (tdf .- b) .* Jprev -c * Jpprev - a = sqrt( dotw(_J, _J) ) - A[n] = 1/a - B[n] = - b / a - C[n] = - c / a - Jprev, Jpprev = _J / a, Jprev - end - - return OrthPolyBasis(pl, tl, pr, tr, A, B, C, collect(tdf), collect(ww)) + pcut::Integer, + tcut::T, + pin::Integer, + tin::T, + tdf::AbstractVector{T}, + ww::AbstractVector{T}=ones(T, length(tdf)) +) where {T<:AbstractFloat} + @assert pcut >= 0 && pin >= 0 + @assert N > 0 + + if tcut < tin + tl, tr = tcut, tin + pl, pr = pcut, pin + else + tl, tr = tin, tcut + pl, pr = pin, pcut + end + + if minimum(tdf) < tl || maximum(tdf) > tr + @warn("OrthPolyBasis: t range outside [tl, tr]") + end + + A = zeros(T, N) + B = zeros(T, N) + C = zeros(T, N) + + # normalise the weights s.t. <1, 1> = 1 + ww = ww ./ sum(ww) + # define inner products + dotw = (f1, f2) -> dot(f1, ww .* f2) + + # start the iteration + _J1 = _fcut_.(pl, tl, pr, tr, tdf) + a = sqrt(dotw(_J1, _J1)) + A[1] = 1 / a + J1 = A[1] * _J1 + + if N > 1 + # a J2 = (t - b) J1 + b = dotw(tdf .* J1, J1) + _J2 = (tdf .- b) .* J1 + a = sqrt(dotw(_J2, _J2)) + A[2] = 1 / a + B[2] = -b / a + J2 = (A[2] * tdf .+ B[2]) .* J1 + + # keep the last two for the 3-term recursion + Jprev = J2 + Jpprev = J1 + end + + for n = 3:N + # a Jn = (t - b) J_{n-1} - c J_{n-2} + b = dotw(tdf .* Jprev, Jprev) + c = dotw(tdf .* Jprev, Jpprev) + _J = (tdf .- b) .* Jprev - c * Jpprev + a = sqrt(dotw(_J, _J)) + A[n] = 1 / a + B[n] = -b / a + C[n] = -c / a + Jprev, Jpprev = _J / a, Jprev + end + + return OrthPolyBasis(pl, tl, pr, tr, A, B, C, collect(tdf), collect(ww)) end -function evaluate(J::OrthPolyBasis, t) - cA = acquire!(J.B_pool, length(J), _valtype(J, t)) - evaluate!(parent(cA), J, t) - return cA +function evaluate(J::OrthPolyBasis, t) + cA = acquire!(J.B_pool, length(J), _valtype(J, t)) + evaluate!(parent(cA), J, t) + return cA end @@ -231,18 +231,22 @@ end evaluate_P1(J::OrthPolyBasis, t) = - J.A[1] * _fcut_(J.pl, J.tl, J.pr, J.tr, t) + J.A[1] * _fcut_(J.pl, J.tl, J.pr, J.tr, t) function evaluate!(P, J::OrthPolyBasis, t; maxn=length(J)) - @assert length(P) >= maxn - P[1] = evaluate_P1(J, t) - if maxn == 1; return P; end - P[2] = (J.A[2] * t + J.B[2]) * P[1] - if maxn == 2; return P; end - @inbounds for n = 3:maxn - P[n] = (J.A[n] * t + J.B[n]) * P[n-1] + J.C[n] * P[n-2] - end - return P + @assert length(P) >= maxn + P[1] = evaluate_P1(J, t) + if maxn == 1 + return P + end + P[2] = (J.A[2] * t + J.B[2]) * P[1] + if maxn == 2 + return P + end + @inbounds for n = 3:maxn + P[n] = (J.A[n] * t + J.B[n]) * P[n-1] + J.C[n] * P[n-2] + end + return P end @@ -280,15 +284,15 @@ using Base: @invokelatest A utility function to generate a jacobi-type basis """ -function discrete_jacobi(N; pcut=0, xcut=1.0, pin=0, xin=-1.0, - Nquad = max(300, 3 * N), - trans = identity) - tcut = @invokelatest trans(xcut) - tin = @invokelatest trans(xin) - tl, tr = minmax(tin, tcut) - dt = (tr - tl) / Nquad - tdf = range(tl + dt/2, tr - dt/2, length=Nquad) - return OrthPolyBasis(N, pcut, tcut, pin, tin, tdf) +function discrete_jacobi(N; pcut=0, xcut=1.0, pin=0, xin=-1.0, + Nquad=max(300, 3 * N), + trans=identity) + tcut = @invokelatest trans(xcut) + tin = @invokelatest trans(xin) + tl, tr = minmax(tin, tcut) + dt = (tr - tl) / Nquad + tdf = range(tl + dt / 2, tr - dt / 2, length=Nquad) + return OrthPolyBasis(N, pcut, tcut, pin, tin, tdf) end @@ -308,14 +312,14 @@ a `TransformPolys` basis with an inner polynomial basis of `OrthPolys` type. * `Nquad = 1000` : number of quadrature points """ function transformed_jacobi(maxdeg::Integer, - trans, - rcut::Real, rin::Real = 0.0; - kwargs...) - J = discrete_jacobi(maxdeg; xcut = rcut, - xin = rin, - pcut = 2, trans=trans, - kwargs...) - return chain(trans, J) + trans, + rcut::Real, rin::Real=0.0; + kwargs...) + J = discrete_jacobi(maxdeg; xcut=rcut, + xin=rin, + pcut=2, trans=trans, + kwargs...) + return chain(trans, J) end diff --git a/src/polynomials/products.jl b/src/polynomials/products.jl index 9e8c843..9a9c6b8 100644 --- a/src/polynomials/products.jl +++ b/src/polynomials/products.jl @@ -5,21 +5,21 @@ using ACEfrictionCore: evaluate # auxiliary datastructure wrapping sparse vectors struct EndlessVector{T} - x::SparseVector{T, Int} + x::SparseVector{T,Int} end -function endless(x::SparseVector{T, Int}; prune=true, tol = 1e-10) where {T} - if prune - x = droptol!(x, tol) - x = x[1:maximum(x.nzind)] - end - return EndlessVector(x) +function endless(x::SparseVector{T,Int}; prune=true, tol=1e-10) where {T} + if prune + x = droptol!(x, tol) + x = x[1:maximum(x.nzind)] + end + return EndlessVector(x) end endless(_x::AbstractVector; kwargs...) = endless(sparse(_x); kwargs...) Base.getindex(x::EndlessVector{T}, i) where {T} = - (0 < i <= length(x.x)) ? x.x[i] : zero(T) + (0 < i <= length(x.x)) ? x.x[i] : zero(T) Base.length(x::EndlessVector) = length(x.x) @@ -28,41 +28,41 @@ Base.length(x::EndlessVector) = length(x.x) # radial basis mutable struct OrthPolyProdCoeffs{T} - basis::OrthPolyBasis{T} - coeffs::Dict{Tuple{Int, Int}, EndlessVector{T}} + basis::OrthPolyBasis{T} + coeffs::AbstractDict{Tuple{Int,Int},EndlessVector{T}} end OrthPolyProdCoeffs(basis::OrthPolyBasis{T}) where {T} = - OrthPolyProdCoeffs(basis, Dict{Tuple{Int, Int}, EndlessVector{T}}()) + OrthPolyProdCoeffs(basis, Dict{Tuple{Int,Int},EndlessVector{T}}()) function (coeffs::OrthPolyProdCoeffs{T})(n1, n2) where {T} - n1, n2 = extrema((n1, n2)) # n1 <= n2 - if n1 <= 0 - return endless(T[]) - end - if !haskey(coeffs.coeffs, (n1, n2)) - coeffs.coeffs[(n1, n2)] = _precompute_prodcoeffs(coeffs, n1, n2) - end - return coeffs.coeffs[(n1, n2)] + n1, n2 = extrema((n1, n2)) # n1 <= n2 + if n1 <= 0 + return endless(T[]) + end + if !haskey(coeffs.coeffs, (n1, n2)) + coeffs.coeffs[(n1, n2)] = _precompute_prodcoeffs(coeffs, n1, n2) + end + return coeffs.coeffs[(n1, n2)] end function _precompute_prodcoeffs(coeffs::OrthPolyProdCoeffs{T}, n1, n2) where {T} - # we want to expand this function in Jn basis - basis = coeffs.basis - f(x) = (J = evaluate(basis, x); J[n1] * J[n2]) - # evaluate basis function with index ν - evalJ(x, ν) = evaluate(basis, x)[ν] - # get the inner product information, normalise the weights s.t. <1, 1> = 1 - # TODO: abstract this out! - tdf = basis.tdf - ww = basis.ww - ww = ww / sum(ww) - dotJ(f1, f2) = dot(f1.(tdf), ww .* f2.(tdf)) + # we want to expand this function in Jn basis + basis = coeffs.basis + f(x) = (J = evaluate(basis, x); J[n1] * J[n2]) + # evaluate basis function with index ν + evalJ(x, ν) = evaluate(basis, x)[ν] + # get the inner product information, normalise the weights s.t. <1, 1> = 1 + # TODO: abstract this out! + tdf = basis.tdf + ww = basis.ww + ww = ww / sum(ww) + dotJ(f1, f2) = dot(f1.(tdf), ww .* f2.(tdf)) - # now we can get the coefficients (note the basis is orthonormal!!) - maxn = (n1 + n2 - 2) + (basis.pl + basis.pr + 1) - P = [ dotJ(f, x -> evalJ(x, ν)) for ν = 1:maxn ] - return endless(P) + # now we can get the coefficients (note the basis is orthonormal!!) + maxn = (n1 + n2 - 2) + (basis.pl + basis.pr + 1) + P = [dotJ(f, x -> evalJ(x, ν)) for ν = 1:maxn] + return endless(P) end diff --git a/src/polynomials/sphericalharmonics.jl b/src/polynomials/sphericalharmonics.jl index cfaf37f..27b042f 100644 --- a/src/polynomials/sphericalharmonics.jl +++ b/src/polynomials/sphericalharmonics.jl @@ -5,15 +5,15 @@ module SphericalHarmonics using StaticArrays, LinearAlgebra -import ACEfrictionCore, ACEbase, ACEfrictionCore.ACEbase024 +import ACEfrictionCore, ACEbase, ACEfrictionCore.ACEbase024 import ACEfrictionCore: evaluate!, - write_dict, read_dict, - ACEBasis, - acquire!, release!, - evaluate + write_dict, read_dict, + ACEBasis, + acquire!, release!, + evaluate -import ACEfrictionCore: VectorPool, ArrayCache +import ACEfrictionCore: VectorPool, ArrayCache export SHBasis @@ -31,23 +31,23 @@ Use `spher2cart` and `cart2spher` to convert between cartesian and spherical coordinates. """ struct SphericalCoords{T} - r::T - cosφ::T - sinφ::T - cosθ::T - sinθ::T + r::T + cosφ::T + sinφ::T + cosθ::T + sinθ::T end -spher2cart(S::SphericalCoords) = S.r * SVector(S.cosφ*S.sinθ, S.sinφ*S.sinθ, S.cosθ) +spher2cart(S::SphericalCoords) = S.r * SVector(S.cosφ * S.sinθ, S.sinφ * S.sinθ, S.cosθ) function cart2spher(R::AbstractVector) - @assert length(R) == 3 - r = norm(R) - φ = atan(R[2], R[1]) - sinφ, cosφ = sincos(φ) - cosθ = R[3] / r - sinθ = sqrt(R[1]^2+R[2]^2) / r - return SphericalCoords(r, cosφ, sinφ, cosθ, sinθ) + @assert length(R) == 3 + r = norm(R) + φ = atan(R[2], R[1]) + sinφ, cosφ = sincos(φ) + cosθ = R[3] / r + sinθ = sqrt(R[1]^2 + R[2]^2) / r + return SphericalCoords(r, cosφ, sinφ, cosθ, sinθ) end SphericalCoords(φ, θ) = SphericalCoords(1.0, cos(φ), sin(φ), cos(θ), sin(θ)) @@ -58,10 +58,10 @@ convert a gradient with respect to spherical coordinates to a gradient with respect to cartesian coordinates """ function dspher_to_dcart(S, f_φ_div_sinθ, f_θ) - r = S.r + eps() - return SVector( - (S.sinφ * f_φ_div_sinθ) + (S.cosφ * S.cosθ * f_θ), - (S.cosφ * f_φ_div_sinθ) + (S.sinφ * S.cosθ * f_θ), - - ( S.sinθ * f_θ) ) / r + r = S.r + eps() + return SVector(-(S.sinφ * f_φ_div_sinθ) + (S.cosφ * S.cosθ * f_θ), + (S.cosφ * f_φ_div_sinθ) + (S.sinφ * S.cosθ * f_θ), + -(S.sinθ * f_θ)) / r end @@ -71,7 +71,7 @@ end # -------------------------------------------------------- """ -`sizeP(maxL):` +`sizeP(maxL):` Return the size of the set of Associated Legendre Polynomials ``P_l^m(x)`` of degree less than or equal to the given maximum degree """ @@ -90,7 +90,7 @@ Return the index into a flat array of Associated Legendre Polynomials `P_l^m` for the given indices `(l,m)`. `P_l^m` are stored in l-major order i.e. `[P(0,0), [P(1,0), P(1,1), P(2,0), ...]`` """ -index_p(l::Integer,m::Integer) = m + div(l*(l+1), 2) + 1 +index_p(l::Integer, m::Integer) = m + div(l * (l + 1), 2) + 1 """ `index_y(l,m):` @@ -99,13 +99,13 @@ for the given indices `(l,m)`. `Y_lm` are stored in l-major order i.e. [Y(0,0), Y(1,-1), Y(1,0), Y(1,1), Y(2,-2), ...] """ -index_y(l::Integer, m::Integer) = m + l + (l*l) + 1 +index_y(l::Integer, m::Integer) = m + l + (l * l) + 1 -function idx2lm(i::Integer) - l = floor(Int, sqrt(i-1) + 1e-10) - m = i - (l + (l*l) + 1) - return l, m -end +function idx2lm(i::Integer) + l = floor(Int, sqrt(i - 1) + 1e-10) + m = i - (l + (l * l) + 1) + return l, m +end # -------------------------------------------------------- @@ -123,86 +123,89 @@ ALPolynomials(maxL::Integer, T::Type=Float64) ``` """ struct ALPolynomials{T} <: ACEBasis - L::Int - A::Vector{T} - B::Vector{T} - B_pool::ArrayCache{T} + L::Int + A::Vector{T} + B::Vector{T} + B_pool::ArrayCache{T} end -ALPolynomials(L::Integer, A::Vector{T}, B::Vector{T}) where {T} = - ALPolynomials(L, A, B, ArrayCache{T}()) +ALPolynomials(L::Integer, A::Vector{T}, B::Vector{T}) where {T} = + ALPolynomials(L, A, B, ArrayCache{T}()) Base.length(alp::ALPolynomials) = sizeP(alp.L) import Base.== -==(B1::ALPolynomials{T}, B2::ALPolynomials{T}) where {T} = - ((B1.L == B2.L) && (B1.A ≈ B2.A) && (B1.B ≈ B2.B)) +==(B1::ALPolynomials{T}, B2::ALPolynomials{T}) where {T} = + ((B1.L == B2.L) && (B1.A ≈ B2.A) && (B1.B ≈ B2.B)) -_valtype(alp::ALPolynomials{T}, x::SphericalCoords{S}) where {T, S} = - promote_type(T, S) +_valtype(alp::ALPolynomials{T}, x::SphericalCoords{S}) where {T,S} = + promote_type(T, S) function ALPolynomials(L::Integer, T::Type=Float64) - # Precompute coefficients ``a_l^m`` and ``b_l^m`` for all l <= L, m <= l - alp = ALPolynomials(L, zeros(T, sizeP(L)), zeros(T, sizeP(L))) - for l in 2:L - ls = l*l - lm1s = (l-1) * (l-1) - for m in 0:(l-2) - ms = m * m - alp.A[index_p(l, m)] = sqrt((4 * ls - 1.0) / (ls - ms)) - alp.B[index_p(l, m)] = -sqrt((lm1s - ms) / (4 * lm1s - 1.0)) - end - end - return alp + # Precompute coefficients ``a_l^m`` and ``b_l^m`` for all l <= L, m <= l + alp = ALPolynomials(L, zeros(T, sizeP(L)), zeros(T, sizeP(L))) + for l in 2:L + ls = l * l + lm1s = (l - 1) * (l - 1) + for m in 0:(l-2) + ms = m * m + alp.A[index_p(l, m)] = sqrt((4 * ls - 1.0) / (ls - ms)) + alp.B[index_p(l, m)] = -sqrt((lm1s - ms) / (4 * lm1s - 1.0)) + end + end + return alp end -function evaluate(alp::ALPolynomials, S::SphericalCoords) - P = acquire!(alp.B_pool, length(alp), _valtype(alp, S)) - evaluate!(parent(P), alp, S) - return P +function evaluate(alp::ALPolynomials, S::SphericalCoords) + P = acquire!(alp.B_pool, length(alp), _valtype(alp, S)) + evaluate!(parent(P), alp, S) + return P end function evaluate!(P, alp::ALPolynomials, S::SphericalCoords) - L = alp.L - A = alp.A - B = alp.B - @assert length(A) >= sizeP(L) - @assert length(B) >= sizeP(L) - @assert length(P) >= sizeP(L) - - temp = sqrt(0.5/π) - P[index_p(0, 0)] = temp - if L == 0; return P; end - - P[index_p(1, 0)] = S.cosθ * sqrt(3) * temp - temp = - sqrt(1.5) * S.sinθ * temp - P[index_p(1, 1)] = temp - - for l in 2:L - il = ((l*(l+1)) ÷ 2) + 1 - ilm1 = il - l - ilm2 = ilm1 - l + 1 - for m in 0:(l-2) - @inbounds P[il+m] = A[il+m] * ( S.cosθ * P[ilm1+m] - + B[il+m] * P[ilm2+m] ) - end - @inbounds P[il+l-1] = S.cosθ * sqrt(2 * (l - 1) + 3) * temp - temp = -sqrt(1.0 + 0.5 / l) * S.sinθ * temp - @inbounds P[il+l] = temp - end - - return P + L = alp.L + A = alp.A + B = alp.B + @assert length(A) >= sizeP(L) + @assert length(B) >= sizeP(L) + @assert length(P) >= sizeP(L) + + temp = sqrt(0.5 / π) + P[index_p(0, 0)] = temp + if L == 0 + return P + end + + P[index_p(1, 0)] = S.cosθ * sqrt(3) * temp + temp = -sqrt(1.5) * S.sinθ * temp + P[index_p(1, 1)] = temp + + for l in 2:L + il = ((l * (l + 1)) ÷ 2) + 1 + ilm1 = il - l + ilm2 = ilm1 - l + 1 + for m in 0:(l-2) + @inbounds P[il+m] = A[il+m] * (S.cosθ * P[ilm1+m] + + + B[il+m] * P[ilm2+m]) + end + @inbounds P[il+l-1] = S.cosθ * sqrt(2 * (l - 1) + 3) * temp + temp = -sqrt(1.0 + 0.5 / l) * S.sinθ * temp + @inbounds P[il+l] = temp + end + + return P end -# this doesn't use the standard name because it doesn't +# this doesn't use the standard name because it doesn't # technically perform the derivative w.r.t. S, but w.r.t. θ # further, P doesn't store P but (P if m = 0) or (P * sinθ if m > 0) -# this is done for numerical stability +# this is done for numerical stability # function _evaluate_ed! removed - derivative functionality has been removed @@ -213,8 +216,8 @@ end # ------------------------------------------------------------------------ """ -`AbstractSHBasis`: This extra abstraction is no longer needed, but there uses -to be a real SH basis and in case this is revived, I am keeping it for now. +`AbstractSHBasis`: This extra abstraction is no longer needed, but there uses +to be a real SH basis and in case this is revived, I am keeping it for now. """ abstract type AbstractSHBasis{T} <: ACEBasis end @@ -222,15 +225,15 @@ abstract type AbstractSHBasis{T} <: ACEBasis end complex spherical harmonics """ struct SHBasis{T} <: AbstractSHBasis{T} - alp::ALPolynomials{T} - B_pool::ArrayCache{Complex{T}} - dB_pool::ArrayCache{SVector{3, Complex{T}}} + alp::ALPolynomials{T} + B_pool::ArrayCache{Complex{T}} + dB_pool::ArrayCache{SVector{3,Complex{T}}} end SHBasis(maxL::Integer, T::Type=Float64) = SHBasis(ALPolynomials(maxL, T)) -SHBasis(alp::ALPolynomials{T}) where {T} = SHBasis(alp, - ArrayCache{Complex{T}}(), ArrayCache{SVector{3, Complex{T}}}()) +SHBasis(alp::ALPolynomials{T}) where {T} = SHBasis(alp, + ArrayCache{Complex{T}}(), ArrayCache{SVector{3,Complex{T}}}()) Base.show(io::IO, SH::SHBasis) = print(io, "SHBasis(L=$(maxL(SH)))") @@ -239,28 +242,28 @@ max L degree for which the alp coefficients have been precomputed """ maxL(sh::AbstractSHBasis) = sh.alp.L -_valtype(sh::SHBasis{T}, x::AbstractVector{S}) where {T, S} = - Complex{promote_type(T, S)} +_valtype(sh::SHBasis{T}, x::AbstractVector{S}) where {T,S} = + Complex{promote_type(T, S)} -_gradtype(sh::SHBasis{T}, x::AbstractVector{S}) where {T, S} = - SVector{3, Complex{promote_type(T, S)}} +_gradtype(sh::SHBasis{T}, x::AbstractVector{S}) where {T,S} = + SVector{3,Complex{promote_type(T, S)}} -function ACEfrictionCore.degree(sh::SHBasis, i::Integer) - l, m = idx2lm(i) - return l +function ACEfrictionCore.degree(sh::SHBasis, i::Integer) + l, m = idx2lm(i) + return l end import Base.== ==(B1::AbstractSHBasis, B2::AbstractSHBasis) = - (B1.alp == B2.alp) && (typeof(B1) == typeof(B2)) + (B1.alp == B2.alp) && (typeof(B1) == typeof(B2)) write_dict(SH::SHBasis{T}) where {T} = - Dict("__id__" => "ACEfrictionCore_SHBasis", - "T" => write_dict(T), - "maxL" => maxL(SH)) + Dict("__id__" => "ACEfrictionCore_SHBasis", + "T" => write_dict(T), + "maxL" => maxL(SH)) -read_dict(::Val{:ACEfrictionCore_SHBasis}, D::Dict) = - SHBasis(D["maxL"], read_dict(D["T"])) +read_dict(::Val{:ACEfrictionCore_SHBasis}, D::AbstractDict) = + SHBasis(D["maxL"], read_dict(D["T"])) Base.length(S::AbstractSHBasis) = sizeY(maxL(S)) @@ -269,19 +272,19 @@ Base.length(S::AbstractSHBasis) = sizeY(maxL(S)) # _evaluate_d! and _evaluate_ed! functions removed - derivative functionality has been removed function ACEfrictionCore.evaluate(SH::SHBasis, R::AbstractVector) - Y = acquire!(SH.B_pool, length(SH), _valtype(SH, R)) - evaluate!(parent(Y), SH, R) - return Y + Y = acquire!(SH.B_pool, length(SH), _valtype(SH, R)) + evaluate!(parent(Y), SH, R) + return Y end function evaluate!(Y, SH::AbstractSHBasis, R::AbstractVector) - @assert length(R) == 3 - L = maxL(SH) - S = cart2spher(R) - P = evaluate(SH.alp, S) - cYlm!(Y, maxL(SH), S, P) - release!(P) - return Y + @assert length(R) == 3 + L = maxL(SH) + S = cart2spher(R) + P = evaluate(SH.alp, S) + cYlm!(Y, maxL(SH), S, P) + release!(P) + return Y end @@ -292,29 +295,29 @@ end evaluate complex spherical harmonics """ function cYlm!(Y, L, S::SphericalCoords, P) - @assert length(P) >= sizeP(L) - @assert length(Y) >= sizeY(L) - @assert abs(S.cosθ) <= 1.0 - - ep = 1 / sqrt(2) + im * 0 - for l = 0:L - Y[index_y(l, 0)] = P[index_p(l, 0)] * ep - end - - sig = 1 - ep_fact = S.cosφ + im * S.sinφ - for m in 1:L - sig *= -1 - ep *= ep_fact # ep = exp(i * m * φ) - em = sig * conj(ep) # ep = ± exp(i * (-m) * φ) - for l in m:L - p = P[index_p(l,m)] - @inbounds Y[index_y(l, -m)] = em * p # (-1)^m * p * exp(-im*m*phi) / sqrt(2) - @inbounds Y[index_y(l, m)] = ep * p # p * exp( im*m*phi) / sqrt(2) - end - end - - return Y + @assert length(P) >= sizeP(L) + @assert length(Y) >= sizeY(L) + @assert abs(S.cosθ) <= 1.0 + + ep = 1 / sqrt(2) + im * 0 + for l = 0:L + Y[index_y(l, 0)] = P[index_p(l, 0)] * ep + end + + sig = 1 + ep_fact = S.cosφ + im * S.sinφ + for m in 1:L + sig *= -1 + ep *= ep_fact # ep = exp(i * m * φ) + em = sig * conj(ep) # ep = ± exp(i * (-m) * φ) + for l in m:L + p = P[index_p(l, m)] + @inbounds Y[index_y(l, -m)] = em * p # (-1)^m * p * exp(-im*m*phi) / sqrt(2) + @inbounds Y[index_y(l, m)] = ep * p # p * exp( im*m*phi) / sqrt(2) + end + end + + return Y end @@ -325,4 +328,3 @@ end end - diff --git a/src/product_1pbasis.jl b/src/product_1pbasis.jl index 9f35268..5a12393 100644 --- a/src/product_1pbasis.jl +++ b/src/product_1pbasis.jl @@ -2,28 +2,28 @@ using NamedTupleTools: namedtuple, merge # -------------- Implementation of Product Basis -struct Product1pBasis{NB, TB <: Tuple} <: OneParticleBasis{Any} - bases::TB - indices::Vector{NTuple{NB, Int}} +struct Product1pBasis{NB,TB<:Tuple} <: OneParticleBasis{Any} + bases::TB + indices::Vector{NTuple{NB,Int}} end function Product1pBasis(bases) - NB = length(bases) - return Product1pBasis(bases, NTuple{NB, Int}[]) + NB = length(bases) + return Product1pBasis(bases, NTuple{NB,Int}[]) end import Base.* *(B1::OneParticleBasis, B2::OneParticleBasis) = - Product1pBasis((B1, B2)) + Product1pBasis((B1, B2)) *(B1::Product1pBasis, B2::OneParticleBasis) = - Product1pBasis((B1.bases..., B2)) + Product1pBasis((B1.bases..., B2)) *(B1::OneParticleBasis, B2::Product1pBasis) = - Product1pBasis((B1, B2.bases...)) + Product1pBasis((B1, B2.bases...)) *(B1::Product1pBasis, B2::Product1pBasis) = - Product1pBasis((B1.bases..., B2.bases...)) + Product1pBasis((B1.bases..., B2.bases...)) *(B1::Product1pBasis, B2::B1pComponent) = - Product1pBasis((B1.bases..., B2)) + Product1pBasis((B1.bases..., B2)) _numb(b::Product1pBasis{NB}) where {NB} = NB @@ -32,255 +32,255 @@ Base.length(basis::Product1pBasis) = length(basis.indices) function Base.show(io::IO, basis::Product1pBasis) - print(io, "Product1pBasis") - print(io, basis.bases) + print(io, "Product1pBasis") + print(io, basis.bases) end -Base.getindex(basis::Product1pBasis, i::Integer) = basis.bases[i] +Base.getindex(basis::Product1pBasis, i::Integer) = basis.bases[i] function Base.getindex(basis::Product1pBasis, label::AbstractString) - inds = findall(getlabel.(basis.bases) .== label) - if length(inds) == 0 - error("label not found amongst 1p basis components") - elseif length(inds) > 1 - error("label not unique amongst 1p basis components") - end - return basis.bases[inds[1]] + inds = findall(getlabel.(basis.bases) .== label) + if length(inds) == 0 + error("label not found amongst 1p basis components") + elseif length(inds) > 1 + error("label not unique amongst 1p basis components") + end + return basis.bases[inds[1]] end # ------------------------- FIO CODES -==(B1::Product1pBasis, B2::Product1pBasis) = - ( all(B1.bases .== B2.bases) && - B1.indices == B2.indices ) +==(B1::Product1pBasis, B2::Product1pBasis) = + (all(B1.bases .== B2.bases) && + B1.indices == B2.indices) -write_dict(B::Product1pBasis) = - Dict("__id__" => "ACEfrictionCore_Product1pBasis", - "bases" => write_dict.(B.bases), - "indices" => B.indices ) +write_dict(B::Product1pBasis) = + Dict("__id__" => "ACEfrictionCore_Product1pBasis", + "bases" => write_dict.(B.bases), + "indices" => B.indices) -function read_dict(::Val{:ACEfrictionCore_Product1pBasis}, D::Dict) - bases = tuple( read_dict.(D["bases"])... ) - indices = [ tuple(v...) for v in D["indices"] ] - return Product1pBasis(bases, indices) +function read_dict(::Val{:ACEfrictionCore_Product1pBasis}, D::AbstractDict) + bases = tuple(read_dict.(D["bases"])...) + indices = [tuple(v...) for v in D["indices"]] + return Product1pBasis(bases, indices) end -# ----------------- evaluation of the basis +# ----------------- evaluation of the basis import Base.Cartesian: @nexprs function _write_A_code(VA, NB) - prodBi_str = "B_1[ϕ[1]]" - for i in 2:NB - prodBi_str *= " * B_$i[ϕ[$i]]" - end - prodBi = Meta.parse(prodBi_str) - if VA == Nothing - getVT = "promote_type(" * prod("eltype(B_$i), " for i = 1:NB) * ")" - getA = Meta.parse("_A = zeros($(getVT), length(basis))") - else - getA = :(_A = A) - end - return prodBi, getA + prodBi_str = "B_1[ϕ[1]]" + for i in 2:NB + prodBi_str *= " * B_$i[ϕ[$i]]" + end + prodBi = Meta.parse(prodBi_str) + if VA == Nothing + getVT = "promote_type(" * prod("eltype(B_$i), " for i = 1:NB) * ")" + getA = Meta.parse("_A = zeros($(getVT), length(basis))") + else + getA = :(_A = A) + end + return prodBi, getA end """ -`add_into_A!` : this is an internal function implementing the main evaluation -for the one-particle basis and possibly add it into the A basis. - -There are two ways to call it. -* Use `A = nothing` for the first argument to allocate the necessary memory to +`add_into_A!` : this is an internal function implementing the main evaluation +for the one-particle basis and possibly add it into the A basis. + +There are two ways to call it. +* Use `A = nothing` for the first argument to allocate the necessary memory to evaluate the 1p basis into it. -* Use `A::Vector{T}` to evaluate the 1p basis and add it into `A` directly -without additional allocation. +* Use `A::Vector{T}` to evaluate the 1p basis and add it into `A` directly +without additional allocation. """ -@generated function add_into_A!(A::VA, basis::Product1pBasis{NB}, X) where {NB, VA} - prodBi, getA = _write_A_code(VA, NB) - quote - # evaluate the 1p basis components - @nexprs $NB i -> begin - bas_i = basis.bases[i] - B_i = evaluate(bas_i, X) - end - # allocate A if necessary or just name _A = A if A is a buffer - $(getA) - # evaluate the 1p product basis functions and add/write into _A - for (iA, ϕ) in enumerate(basis.indices) - @inbounds _A[iA] += $prodBi - end - # release the memory allocated by the 1p basis components they normally - # use preallocated chache to avoid too many small allocations. - @nexprs $NB i -> release!(B_i) - return _A - end +@generated function add_into_A!(A::VA, basis::Product1pBasis{NB}, X) where {NB,VA} + prodBi, getA = _write_A_code(VA, NB) + quote + # evaluate the 1p basis components + @nexprs $NB i -> begin + bas_i = basis.bases[i] + B_i = evaluate(bas_i, X) + end + # allocate A if necessary or just name _A = A if A is a buffer + $(getA) + # evaluate the 1p product basis functions and add/write into _A + for (iA, ϕ) in enumerate(basis.indices) + @inbounds _A[iA] += $prodBi + end + # release the memory allocated by the 1p basis components they normally + # use preallocated chache to avoid too many small allocations. + @nexprs $NB i -> release!(B_i) + return _A + end end -evaluate(basis::Product1pBasis, X::AbstractState) = - add_into_A!(nothing, basis, X) +evaluate(basis::Product1pBasis, X::AbstractState) = + add_into_A!(nothing, basis, X) function evaluate(basis::Product1pBasis, cfg::UConfig) - @assert length(cfg) > 0 "Product1pBasis can only be evaluated with non-empty configurations" - # evaluate the first item "manually", then so we know the output types - # but then write directly into the allocated array to avoid additional - # allocations. - A = evaluate(basis, first(cfg)) - for (i, X) in enumerate(cfg) - i == 1 && continue; - add_into_A!(A, basis, X) - end - return A -end + @assert length(cfg) > 0 "Product1pBasis can only be evaluated with non-empty configurations" + # evaluate the first item "manually", then so we know the output types + # but then write directly into the allocated array to avoid additional + # allocations. + A = evaluate(basis, first(cfg)) + for (i, X) in enumerate(cfg) + i == 1 && continue + add_into_A!(A, basis, X) + end + return A +end function evaluate!(A::AbstractVector, basis::Product1pBasis, X::AbstractState) - fill!(A, zero(eltype(A))) - add_into_A!(A, basis, X) - return A + fill!(A, zero(eltype(A))) + add_into_A!(A, basis, X) + return A end function evaluate!(A, basis::Product1pBasis, cfg::UConfig) - fill!(A, zero(eltype(A))) - for X in cfg - add_into_A!(A, basis, X) - end - return A + fill!(A, zero(eltype(A))) + for X in cfg + add_into_A!(A, basis, X) + end + return A end -# ------------- Partial derivative functionality +# ------------- Partial derivative functionality -_check_args_is_sym() = true +_check_args_is_sym() = true _check_args_is_sym(::Symbol) = true # ---------------------------------------- -_symbols_prod(bases) = tuple(union( symbols.(bases)... )...) +_symbols_prod(bases) = tuple(union(symbols.(bases)...)...) symbols(basis::Product1pBasis) = _symbols_prod(basis.bases) function indexrange(basis::Product1pBasis) - allsyms = tuple(symbols(basis)...) - rg = Dict{Symbol, Vector{Any}}([ sym => [] for sym in allsyms]...) - for b in basis.bases - rgb = indexrange(b) - for sym in allsyms - if haskey(rgb, sym) - rg[sym] = union(rg[sym], rgb[sym]) - end - end - end - # HACK: fix the m range based on the maximal l-range - # this needs to be suitably generalised if we have multiple - # (l, m) pairs, e.g. (l1, m1), (l2, m2) - if haskey(rg, :m) - maxl = maximum(rg[:l]) - rg[:m] = collect(-maxl:maxl) - end - - # convert the range into a named tuple so that we remember the order!! - return NamedTuple{allsyms}(ntuple(i -> rg[allsyms[i]], length(allsyms))) + allsyms = tuple(symbols(basis)...) + rg = Dict{Symbol,Vector{Any}}([sym => [] for sym in allsyms]...) + for b in basis.bases + rgb = indexrange(b) + for sym in allsyms + if haskey(rgb, sym) + rg[sym] = union(rg[sym], rgb[sym]) + end + end + end + # HACK: fix the m range based on the maximal l-range + # this needs to be suitably generalised if we have multiple + # (l, m) pairs, e.g. (l1, m1), (l2, m2) + if haskey(rg, :m) + maxl = maximum(rg[:l]) + rg[:m] = collect(-maxl:maxl) + end + + # convert the range into a named tuple so that we remember the order!! + return NamedTuple{allsyms}(ntuple(i -> rg[allsyms[i]], length(allsyms))) end isadmissible(b, basis::Product1pBasis) = all(isadmissible.(Ref(b), basis.bases)) function set_spec!(basis::Product1pBasis{NB}, spec) where {NB} - empty!(basis.indices) - for b in spec - inds = ntuple(i -> get_index(basis.bases[i], b), NB) - push!(basis.indices, inds) - end - return basis + empty!(basis.indices) + for b in spec + inds = ntuple(i -> get_index(basis.bases[i], b), NB) + push!(basis.indices, inds) + end + return basis end -get_spec(basis::Product1pBasis) = [ get_spec(basis, i) for i = 1:length(basis) ] +get_spec(basis::Product1pBasis) = [get_spec(basis, i) for i = 1:length(basis)] -function get_spec(basis::Product1pBasis, i::Integer) - inds = basis.indices[i] - specs = get_spec.(basis.bases, inds) - # TODO: here we should check that we are only merging compatible tuples, - # e.g. (n = 5, l = 2), (l = 2, m = -1) is ok - # but (n = 5, l = 2), (l = 3, m = -1) is forbidden! - return merge(specs...) +function get_spec(basis::Product1pBasis, i::Integer) + inds = basis.indices[i] + specs = get_spec.(basis.bases, inds) + # TODO: here we should check that we are only merging compatible tuples, + # e.g. (n = 5, l = 2), (l = 2, m = -1) is ok + # but (n = 5, l = 2), (l = 3, m = -1) is forbidden! + return merge(specs...) end -degree(b, basis::Product1pBasis) = sum( degree(b, B) for B in basis.bases ) +degree(b, basis::Product1pBasis) = sum(degree(b, B) for B in basis.bases) -degree(b::NamedTuple, basis::Product1pBasis, weight::Dict) = - sum( degree(b, B, weight) for B in basis.bases ) +degree(b::NamedTuple, basis::Product1pBasis, weight::AbstractDict) = + sum(degree(b, B, weight) for B in basis.bases) # TODO: this looks like a horrible hack ... function rand_radial(basis::Product1pBasis) - for B in basis.bases - if B isa ScalarACEBasis - return rand_radial(B) - end - end - return nothing + for B in basis.bases + if B isa ScalarACEBasis + return rand_radial(B) + end + end + return nothing end -# -------------- sparsification - -function sparsify!(basis1p::Product1pBasis, keep::AbstractVector{<: NamedTuple}) - # spec, keep, new_spec will be lists of named tuples, - # e.g. [ (n = , l = , m = ), ... ] - spec = get_spec(basis1p) - new_spec = eltype(spec)[] - new_inds = Vector{Int}(undef, length(spec)) - for (ib, b) in enumerate(spec) - if b in keep - push!(new_spec, b) - new_inds[ib] = length(new_spec) - end - end - - # now we need to recompute the indices array, this can be easily done via - # set_spec!(basis::Product1pBasis{NB}, spec), but before we do that - # we should sparsify the basis components as well - # .... but it is not so clear that his is a good idea, maybe the - # 1p basis components should just remain frozen??? - # => turn this off for now - # TODO - return to this point?!?!? - # for bas_i in basis1p.bases - # _sparsify_component!(bas_i, new_spec) - # end - - # finally fix the basis1pspec internally: - set_spec!(basis1p, new_spec) - - # return the old to new index mapping so that the pibasis can fix itself. - return basis1p, new_inds +# -------------- sparsification + +function sparsify!(basis1p::Product1pBasis, keep::AbstractVector{<:NamedTuple}) + # spec, keep, new_spec will be lists of named tuples, + # e.g. [ (n = , l = , m = ), ... ] + spec = get_spec(basis1p) + new_spec = eltype(spec)[] + new_inds = Vector{Int}(undef, length(spec)) + for (ib, b) in enumerate(spec) + if b in keep + push!(new_spec, b) + new_inds[ib] = length(new_spec) + end + end + + # now we need to recompute the indices array, this can be easily done via + # set_spec!(basis::Product1pBasis{NB}, spec), but before we do that + # we should sparsify the basis components as well + # .... but it is not so clear that his is a good idea, maybe the + # 1p basis components should just remain frozen??? + # => turn this off for now + # TODO - return to this point?!?!? + # for bas_i in basis1p.bases + # _sparsify_component!(bas_i, new_spec) + # end + + # finally fix the basis1pspec internally: + set_spec!(basis1p, new_spec) + + # return the old to new index mapping so that the pibasis can fix itself. + return basis1p, new_inds end -using NamedTupleTools: select +using NamedTupleTools: select # """ -# this performs some generic work to sparsify a 1p-basis component. -# but the actual sparsificatin happens in the individual basis implementations +# this performs some generic work to sparsify a 1p-basis component. +# but the actual sparsificatin happens in the individual basis implementations # """ # function _sparsify_component!(basis1p, keep) -# # if basis1p has no symbols (e.g. a multiplier) then it means it must +# # if basis1p has no symbols (e.g. a multiplier) then it means it must # # be a one-component basis, so there is nothing to sparsify. # syms = symbols(basis1p) # if isempty(syms) # return basis1p # end -# # get rid of all info we don't need +# # get rid of all info we don't need # keep1 = unique( select.(keep, Ref(syms)) ) -# # double-check that keep1 is compatible -# spec = get_spec(basis1p) +# # double-check that keep1 is compatible +# spec = get_spec(basis1p) # @assert all(b in spec for b in keep1) -# # now get the basis spec and get the list of indices to keep +# # now get the basis spec and get the list of indices to keep # if length(keep1) < length(spec) # # Ikeep = findall( [b in keep1 for b in spec] ) # # sparsify!(basis1p, Ikeep) # sparsify!(basis1p, keep1) -# end -# return basis1p +# end +# return basis1p # end @@ -288,31 +288,31 @@ using NamedTupleTools: select # import ChainRules: rrule, NoTangent, ZeroTangent -# _evaluate_bases(basis::Product1pBasis{NB}, X::AbstractState) where {NB} = +# _evaluate_bases(basis::Product1pBasis{NB}, X::AbstractState) where {NB} = # ntuple(i -> evaluate(basis.bases[i], X), NB) -# _evaluate_A(basis::Product1pBasis{NB}, BB) where {NB} = +# _evaluate_A(basis::Product1pBasis{NB}, BB) where {NB} = # [ prod(BB[i][ϕ[i]] for i = 1:NB) for ϕ in basis.indices ] -# evaluate(basis::Product1pBasis, X::AbstractState) = -# _evaluate_A(basis, _evaluate_bases(basis, X)) +# evaluate(basis::Product1pBasis, X::AbstractState) = +# _evaluate_A(basis, _evaluate_bases(basis, X)) -# function _rrule_evaluate(basis::Product1pBasis{NB}, X::AbstractState, -# w::AbstractVector{<: Number}, +# function _rrule_evaluate(basis::Product1pBasis{NB}, X::AbstractState, +# w::AbstractVector{<: Number}, # BB = _evaluate_bases(basis, X)) where {NB} # VT = promote_type(valtype(basis, X), eltype(w)) # # dB = evaluate_d(basis, X) -# # return sum( (real(w) * real(db) + imag(w) * imag(db)) +# # return sum( (real(w) * real(db) + imag(w) * imag(db)) # # for (w, db) in zip(w, dB) ) -# # Compute the differentials for the individual sub-bases -# Wsub = ntuple(i -> zeros(VT, length(BB[i])), NB) +# # Compute the differentials for the individual sub-bases +# Wsub = ntuple(i -> zeros(VT, length(BB[i])), NB) # for (ivv, vv) in enumerate(basis.indices) -# for t = 1:NB +# for t = 1:NB # _A = one(VT) -# for s = 1:NB -# if s != t +# for s = 1:NB +# if s != t # _A *= BB[s][vv[s]] # end # end @@ -320,8 +320,8 @@ using NamedTupleTools: select # end # end -# # now these can be propagated into the inner basis -# # -> type instab to be fixed here +# # now these can be propagated into the inner basis +# # -> type instab to be fixed here # g = sum( _rrule_evaluate(basis.bases[t], X, Wsub[t] ) # for t = 1:NB ) # return g @@ -330,12 +330,12 @@ using NamedTupleTools: select # function rrule(::typeof(evaluate), basis::Product1pBasis, X::AbstractState) # BB = _evaluate_bases(basis, X) # A = _evaluate_A(basis, BB) -# return A, +# return A, # w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w, BB)) # end -# function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState, +# function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState, # w::AbstractVector{<: Number}) # x = _val(X, basis) # a = _rrule_evaluate(basis.P, x, w) diff --git a/src/properties.jl b/src/properties.jl index b8ddb7f..043a204 100644 --- a/src/properties.jl +++ b/src/properties.jl @@ -4,17 +4,17 @@ import LinearAlgebra: norm, promote_leaf_eltypes _basetype(φ::AbstractProperty) = Base.typename(typeof(φ)).wrapper -# TODO: rewrite all this with meta-programming +# TODO: rewrite all this with meta-programming -@inline +(φ1::AbstractProperty, φ2::AbstractProperty) = (_basetype(φ1))( φ1.val + φ2.val ) -@inline -(φ1::AbstractProperty, φ2::AbstractProperty) = (_basetype(φ1))( φ1.val - φ2.val ) -@inline -(φ1::AbstractProperty) = (_basetype(φ1))( -φ1.val ) -@inline *(φ1::AbstractProperty, φ2::AbstractProperty) = (_basetype(φ1))( φ1.val * φ2.val ) +@inline +(φ1::AbstractProperty, φ2::AbstractProperty) = (_basetype(φ1))(φ1.val + φ2.val) +@inline -(φ1::AbstractProperty, φ2::AbstractProperty) = (_basetype(φ1))(φ1.val - φ2.val) +@inline -(φ1::AbstractProperty) = (_basetype(φ1))(-φ1.val) +@inline *(φ1::AbstractProperty, φ2::AbstractProperty) = (_basetype(φ1))(φ1.val * φ2.val) -@inline *(a::Union{Number, AbstractMatrix}, φ::AbstractProperty) = (_basetype(φ))(a * φ.val) -@inline *(φ::AbstractProperty, a::Union{Number, AbstractMatrix}) = (_basetype(φ))(φ.val * a) +@inline *(a::Union{Number,AbstractMatrix}, φ::AbstractProperty) = (_basetype(φ))(a * φ.val) +@inline *(φ::AbstractProperty, a::Union{Number,AbstractMatrix}) = (_basetype(φ))(φ.val * a) -*(a::AbstractVector{<: Number}, φ::AbstractProperty) = a .* Ref(φ) +*(a::AbstractVector{<:Number}, φ::AbstractProperty) = a .* Ref(φ) Base.isapprox(φ1::AbstractProperty, φ2::AbstractProperty) = isapprox(φ1.val, φ2.val) @@ -22,52 +22,52 @@ Base.isapprox(φ1::AbstractProperty, φ2::AbstractProperty) = isapprox(φ1.val, @inline Base.length(φ::AbstractProperty) = length(φ.val) @inline Base.size(φ::AbstractProperty) = size(φ.val) @inline Base.zero(φ::AbstractProperty) = (_basetype(φ))(zero(φ.val)) -@inline Base.zero(::Type{T}) where {T <: AbstractProperty} = zero(T()) +@inline Base.zero(::Type{T}) where {T<:AbstractProperty} = zero(T()) -promote_leaf_eltypes(φ::T) where {T <: AbstractProperty} = promote_leaf_eltypes(φ.val) +promote_leaf_eltypes(φ::T) where {T<:AbstractProperty} = promote_leaf_eltypes(φ.val) -Base.convert(T::Type{TP}, φ::TP) where {TP <: AbstractProperty} = φ +Base.convert(T::Type{TP}, φ::TP) where {TP<:AbstractProperty} = φ Base.convert(T::Type, φ::AbstractProperty) = convert(T, φ.val) Base.convert(T::Type{Any}, φ::AbstractProperty) = φ -# TODO: this was a hack; not clear this is a good idea ... -# remove during next cleanup and see what happens? +# TODO: this was a hack; not clear this is a good idea ... +# remove during next cleanup and see what happens? Base.iterate(φ::AbstractProperty) = φ, nothing Base.iterate(φ::AbstractProperty, ::Nothing) = nothing """ -`coco_o_daa` : implements a tensor product between a coupling coefficient -(usually an `AbstractProperty`) and a gradient (usually a `DState`). +`coco_o_daa` : implements a tensor product between a coupling coefficient +(usually an `AbstractProperty`) and a gradient (usually a `DState`). """ -function coco_o_daa(φ::AbstractProperty, b::TX) where {TX <: XState} - SYMS = _syms(TX) - vals = ntuple( i -> coco_o_daa(φ.val, _x(b)[SYMS[i]]), length(SYMS) ) - return TX( NamedTuple{SYMS}(vals) ) +function coco_o_daa(φ::AbstractProperty, b::TX) where {TX<:XState} + SYMS = _syms(TX) + vals = ntuple(i -> coco_o_daa(φ.val, _x(b)[SYMS[i]]), length(SYMS)) + return TX(NamedTuple{SYMS}(vals)) end -coco_o_daa(cc::SVector{N, <: AbstractProperty}, b::TX) where {N, TX <: XState} = - SVector( ntuple(i -> coco_o_daa(cc[i], b), N) ) +coco_o_daa(cc::SVector{N,<:AbstractProperty}, b::TX) where {N,TX<:XState} = + SVector(ntuple(i -> coco_o_daa(cc[i], b), N)) coco_o_daa(cc::Number, b::Number) = cc * b coco_o_daa(cc::Number, b::SVector) = cc * b coco_o_daa(cc::SVector, b::SVector) = cc * transpose(b) coco_o_daa(cc::SMatrix{N1,N2}, b::SVector{N3}) where {N1,N2,N3} = - reshape(cc[:] * transpose(b), Size(N1, N2, N3)) + reshape(cc[:] * transpose(b), Size(N1, N2, N3)) coco_o_daa(cc::SArray{Tuple{N1,N2,N3}}, b::SVector{N4}) where {N1,N2,N3,N4} = - reshape(cc[:] * transpose(b), Size(N1, N2, N3, N4)) + reshape(cc[:] * transpose(b), Size(N1, N2, N3, N4)) -# TODO: is this needed or can it be removed? +# TODO: is this needed or can it be removed? # maybe it should also be allowed for a DState? -# it is also more of an ⊗ ... revisit during next cleanup +# it is also more of an ⊗ ... revisit during next cleanup *(φ::AbstractProperty, b::AbstractState) = coco_o_daa(φ, b) -# default behaviour - overwritten e.g. for EuclideanVector where +# default behaviour - overwritten e.g. for EuclideanVector where # cocos are always complex independently of whether phi is real. -coco_type(T::Type{<: AbstractProperty}) = T +coco_type(T::Type{<:AbstractProperty}) = T """ @@ -75,7 +75,7 @@ coco_type(T::Type{<: AbstractProperty}) = T an invariant scalar. """ struct Invariant{T} <: AbstractProperty - val::T + val::T end Base.show(io::IO, φ::Invariant) = print(io, "i($(φ.val))") @@ -83,19 +83,19 @@ Base.show(io::IO, φ::Invariant) = print(io, "i($(φ.val))") Base.one(φ::Invariant{T}) where {T} = Invariant(one(T)) Base.one(::Type{Invariant{T}}) where {T} = Invariant(one(T)) -isrealB(::Invariant{<: Real}) = true -isrealB(::Invariant{<: Complex}) = false -isrealAA(::Invariant{<: Real}) = true -isrealAA(::Invariant{<: Complex}) = false +isrealB(::Invariant{<:Real}) = true +isrealB(::Invariant{<:Complex}) = false +isrealAA(::Invariant{<:Real}) = true +isrealAA(::Invariant{<:Complex}) = false -Invariant{T}() where {T <: Number} = Invariant{T}(zero(T)) +Invariant{T}() where {T<:Number} = Invariant{T}(zero(T)) -Invariant(T::DataType = Float64) = Invariant{T}() +Invariant(T::DataType=Float64) = Invariant{T}() real(φ::Invariant) = Invariant(real(φ.val)) complex(φ::Invariant) = Invariant(complex(φ.val)) complex(::Type{Invariant{T}}) where {T} = Invariant{complex(T)} -complex(φ::AbstractVector{<: Invariant}) = complex.(φ) +complex(φ::AbstractVector{<:Invariant}) = complex.(φ) +(φ::Invariant, x::Number) = Invariant(φ.val + x) +(x::Number, φ::Invariant) = Invariant(φ.val + x) @@ -103,48 +103,48 @@ Base.convert(::Type{Invariant{T}}, x::Number) where {T} = Invariant(convert(T, x *(φ1::Invariant, φ2::Invariant) = Invariant(φ1.val * φ2.val) -write_dict(φ::Invariant{T}) where {T} = - Dict("__id__" => "ACEfrictionCore_Invariant", +write_dict(φ::Invariant{T}) where {T} = + Dict("__id__" => "ACEfrictionCore_Invariant", "val" => φ.val, - "T" => write_dict(T) ) - -read_dict(::Val{:ACEfrictionCore_Invariant}, D::Dict) = - Invariant{read_dict(D["T"])}(D["val"]) - - -function filter(φ::Invariant, grp::O3, b::Array) - if length(b) <= 1 - return true - end - suml = sum( getl(grp, bi) for bi in b ) - if haskey(b[1], msym(grp)) # depends on context whether m come along? - summ = sum( getm(grp, bi) for bi in b ) - return iseven(suml) && iszero(summ) - end - return iseven(suml) + "T" => write_dict(T)) + +read_dict(::Val{:ACEfrictionCore_Invariant}, D::AbstractDict) = + Invariant{read_dict(D["T"])}(D["val"]) + + +function filter(φ::Invariant, grp::O3, b::Array) + if length(b) <= 1 + return true + end + suml = sum(getl(grp, bi) for bi in b) + if haskey(b[1], msym(grp)) # depends on context whether m come along? + summ = sum(getm(grp, bi) for bi in b) + return iseven(suml) && iszero(summ) + end + return iseven(suml) end -filter(φ::Invariant, grp::O3O3, b::Array) = - filter(φ, grp.G1, b) && filter(φ, grp.G2, b) - +filter(φ::Invariant, grp::O3O3, b::Array) = + filter(φ, grp.G1, b) && filter(φ, grp.G2, b) + filter(φ::AbstractProperty, grp::NoSym, b::Array) = true rot3Dcoeffs(::Invariant, T=Float64) = Rot3DCoeffs(T) *(φ::Invariant, dAA::SVector) = φ.val * dAA -coco_init(::Invariant{T}) where {T} = [ Invariant(complex(T(1)));; ] +coco_init(::Invariant{T}) where {T} = [Invariant(complex(T(1)));;] coco_init(::Invariant, l, m, μ, T, A) = ( - l == m == μ == 0 ? Invariant(T(1)) : Invariant(T(0)) ) + l == m == μ == 0 ? Invariant(T(1)) : Invariant(T(0))) coco_zeros(::Invariant, ll, mm, kk, T, A) = Invariant(zero(T)) coco_filter(::Invariant, ll, mm) = - iseven(sum(ll)) && (sum(mm) == 0) + iseven(sum(ll)) && (sum(mm) == 0) coco_filter(::Invariant, ll, mm, kk) = - iseven(sum(ll)) && (sum(mm) == sum(kk) == 0) + iseven(sum(ll)) && (sum(mm) == sum(kk) == 0) coco_dot(u1::Invariant, u2::Invariant) = u1.val * u2.val @@ -160,63 +160,63 @@ $O(3)$ as where $\cdot$ denotes the standard matrix-vector product. """ struct EuclideanVector{T} <: AbstractProperty - val::SVector{3, T} + val::SVector{3,T} end -Base.show(io::IO, φ::EuclideanVector) = - print(io, "e[$(φ.val[1]), $(φ.val[2]), $(φ.val[3])]") +Base.show(io::IO, φ::EuclideanVector) = + print(io, "e[$(φ.val[1]), $(φ.val[2]), $(φ.val[3])]") real(φ::EuclideanVector) = EuclideanVector(real.(φ.val)) complex(φ::EuclideanVector) = EuclideanVector(complex(φ.val)) complex(::Type{EuclideanVector{T}}) where {T} = EuclideanVector{complex(T)} +(x::SVector{3}, y::EuclideanVector) = EuclideanVector(x + y.val) -Base.convert(::Type{SVector{3, T}}, φ::EuclideanVector) where {T} = convert(SVector{3, T}, φ.val) +Base.convert(::Type{SVector{3,T}}, φ::EuclideanVector) where {T} = convert(SVector{3,T}, φ.val) isrealB(::EuclideanVector{T}) where {T} = (T == real(T)) isrealAA(::EuclideanVector) = false -# CO: removed this, is it needed? -# Base.getindex(φ::EuclideanVector, i::Integer) = φ.val[i] +# CO: removed this, is it needed? +# Base.getindex(φ::EuclideanVector, i::Integer) = φ.val[i] -EuclideanVector{T}() where {T <: Number} = EuclideanVector{T}(zero(SVector{3, T})) +EuclideanVector{T}() where {T<:Number} = EuclideanVector{T}(zero(SVector{3,T})) EuclideanVector(T::DataType=Float64) = EuclideanVector{T}() function filter(φ::EuclideanVector, grp::O3, bb::Array) - if length(bb) == 0 # no zero-correlations allowed - return false - end - if length(bb) == 1 #MS: Not sure if this should be here - # CO: good question - need to investigate - return true - end - suml = sum( getl(grp, bi) for bi in bb ) - if haskey(bb[1], msym(grp)) # depends on context whether m come along? - summ = sum( getm(grp, bi) for bi in bb ) - return isodd(suml) && abs(summ) <= 1 - end - return isodd(suml) + if length(bb) == 0 # no zero-correlations allowed + return false + end + if length(bb) == 1 #MS: Not sure if this should be here + # CO: good question - need to investigate + return true + end + suml = sum(getl(grp, bi) for bi in bb) + if haskey(bb[1], msym(grp)) # depends on context whether m come along? + summ = sum(getm(grp, bi) for bi in bb) + return isodd(suml) && abs(summ) <= 1 + end + return isodd(suml) end -rot3Dcoeffs(::EuclideanVector,T=Float64) = Rot3DCoeffsEquiv{T,1}(Dict[], ClebschGordan(T)) +rot3Dcoeffs(::EuclideanVector, T=Float64) = Rot3DCoeffsEquiv{T,1}(Dict[], ClebschGordan(T)) write_dict(φ::EuclideanVector{T}) where {T} = - Dict("__id__" => "ACEfrictionCore_EuclideanVector", - "val" => write_dict(Vector(φ.val)) ) + Dict("__id__" => "ACEfrictionCore_EuclideanVector", + "val" => write_dict(Vector(φ.val))) -function read_dict(::Val{:ACEfrictionCore_EuclideanVector}, D::Dict) - return EuclideanVector(SVector{3}(read_dict(D["val"]))) +function read_dict(::Val{:ACEfrictionCore_EuclideanVector}, D::AbstractDict) + return EuclideanVector(SVector{3}(read_dict(D["val"]))) end # differentiation - cf #27 # *(φ::EuclideanVector, dAA::SVector) = φ.val * dAA' -*(prop::EuclideanVector, c::SVector{N, T}) where {T<:Number,N} = SVector{N}(prop*c[i] for i=1:N) +*(prop::EuclideanVector, c::SVector{N,T}) where {T<:Number,N} = SVector{N}(prop * c[i] for i = 1:N) coco_init(phi::EuclideanVector{CT}, l, m, μ, T, A) where {CT<:Real} = ( - (l == 1 && abs(m) <= 1 && abs(μ) <= 1) - ? [EuclideanVector(conj(crmatrices[(l=l,m=-m,mu=-μ,i=i)])) for i=1:3] - : coco_zeros(phi, l, m, μ, T, A) ) + (l == 1 && abs(m) <= 1 && abs(μ) <= 1) + ? [EuclideanVector(conj(crmatrices[(l=l, m=-m, mu=-μ, i=i)])) for i = 1:3] + : coco_zeros(phi, l, m, μ, T, A)) #coco_init(phi::EuclideanVector{CT}, l, m, μ, T, A) where {CT<:Real} = ( # (l == 1 && abs(m) <= 1 && abs(μ) <= 1) @@ -224,9 +224,9 @@ coco_init(phi::EuclideanVector{CT}, l, m, μ, T, A) where {CT<:Real} = ( # : coco_zeros(phi, l, m, μ, T, A) ) # coco_init(phi::EuclideanVector{CT} -# this is not needed, since the EuclideanVector should never give us -# a constant basis anyhow. Still ... this could become a problem if we -# ever want to artificially increase the AA basis in order to get some +# this is not needed, since the EuclideanVector should never give us +# a constant basis anyhow. Still ... this could become a problem if we +# ever want to artificially increase the AA basis in order to get some # savings elsewhere. Maybe need to revisit this... coco_type(φ::EuclideanVector) = typeof(complex(φ)) @@ -235,12 +235,12 @@ coco_type(::Type{EuclideanVector{T}}) where {T} = EuclideanVector{complex(T)} coco_zeros(φ::EuclideanVector, ll, mm, kk, T, A) = zeros(typeof(complex(φ)), 3) coco_filter(::EuclideanVector, ll, mm) = - isodd(sum(ll)) && (abs(sum(mm)) <= 1) + isodd(sum(ll)) && (abs(sum(mm)) <= 1) coco_filter(::EuclideanVector, ll, mm, kk) = - abs(sum(mm)) <= 1 && - abs(sum(kk)) <= 1 && - isodd(sum(ll)) + abs(sum(mm)) <= 1 && + abs(sum(kk)) <= 1 && + isodd(sum(ll)) coco_dot(u1::EuclideanVector, u2::EuclideanVector) = dot(u1.val, u2.val) @@ -250,110 +250,110 @@ include("eucl/cov_coeffs_dict.jl") abstract type AbstractEuclideanMatrix{T} <: AbstractProperty end -real(φ::E) where { E <: AbstractEuclideanMatrix} = E(real.(φ.val)) +real(φ::E) where {E<:AbstractEuclideanMatrix} = E(real.(φ.val)) -+(x::SMatrix{3}, y::E) where { E <: AbstractEuclideanMatrix} = E(x + y.val) ++(x::SMatrix{3}, y::E) where {E<:AbstractEuclideanMatrix} = E(x + y.val) -Base.convert(::Type{SMatrix{3, 3, T, 9}}, φ::AbstractEuclideanMatrix) where {T} = convert(SMatrix{3, 3, T, 9}, φ.val) +Base.convert(::Type{SMatrix{3,3,T,9}}, φ::AbstractEuclideanMatrix) where {T} = convert(SMatrix{3,3,T,9}, φ.val) -isrealAA(::AbstractEuclideanMatrix) = false +isrealAA(::AbstractEuclideanMatrix) = false function filter(::AbstractEuclideanMatrix, grp::O3, bb::Array) - if length(bb) == 0 # no zero-correlations allowed - return false - end - if length(bb) == 1 #MS: Not sure if this should be here - return true - end - suml = sum( getl(grp, bi) for bi in bb ) - if haskey(bb[1], msym(grp)) # depends on context whether m come along? - summ = sum( getm(grp, bi) for bi in bb ) - return iseven(suml) && abs(summ) <= 2 - end - return iseven(suml) + if length(bb) == 0 # no zero-correlations allowed + return false + end + if length(bb) == 1 #MS: Not sure if this should be here + return true + end + suml = sum(getl(grp, bi) for bi in bb) + if haskey(bb[1], msym(grp)) # depends on context whether m come along? + summ = sum(getm(grp, bi) for bi in bb) + return iseven(suml) && abs(summ) <= 2 + end + return iseven(suml) end -rot3Dcoeffs(::AbstractEuclideanMatrix,T=Float64) = Rot3DCoeffsEquiv{T,1}(Dict[], ClebschGordan(T)) +rot3Dcoeffs(::AbstractEuclideanMatrix, T=Float64) = Rot3DCoeffsEquiv{T,1}(Dict[], ClebschGordan(T)) coco_type(φ::AbstractEuclideanMatrix) = typeof(complex(φ)) -coco_zeros(φ::AbstractEuclideanMatrix, ll, mm, kk, T, A) = zeros(typeof(complex(φ)), 9) +coco_zeros(φ::AbstractEuclideanMatrix, ll, mm, kk, T, A) = zeros(typeof(complex(φ)), 9) coco_filter(::AbstractEuclideanMatrix, ll, mm) = iseven(sum(ll)) && (abs(sum(mm)) <= 2) -coco_filter(::AbstractEuclideanMatrix, ll, mm, kk) = abs(sum(mm)) <= 2 && - abs(sum(kk)) <= 2 && - iseven(sum(ll)) -coco_dot(u1::AbstractEuclideanMatrix, u2::AbstractEuclideanMatrix) = sum(transpose(conj.( u1.val)) * u2.val) +coco_filter(::AbstractEuclideanMatrix, ll, mm, kk) = abs(sum(mm)) <= 2 && + abs(sum(kk)) <= 2 && + iseven(sum(ll)) +coco_dot(u1::AbstractEuclideanMatrix, u2::AbstractEuclideanMatrix) = sum(transpose(conj.(u1.val)) * u2.val) -*(prop::ACEfrictionCore.AbstractEuclideanMatrix, c::SVector{N, T}) where {T<:Number,N} = SVector{N}(prop*c[i] for i=1:N) +*(prop::ACEfrictionCore.AbstractEuclideanMatrix, c::SVector{N,T}) where {T<:Number,N} = SVector{N}(prop * c[i] for i = 1:N) function Base.show(io::IO, φ::AbstractEuclideanMatrix) - # println(io, "3x3 $(typeof(φ)):") - println(io, "$(_type_marker(φ))[ $(φ.val[1,1]), $(φ.val[1,2]), $(φ.val[1,3]);") - println(io, " $(φ.val[2,1]), $(φ.val[2,2]), $(φ.val[2,3]);") - print(io, " $(φ.val[3,1]), $(φ.val[3,2]), $(φ.val[3,3]) ]") + # println(io, "3x3 $(typeof(φ)):") + println(io, "$(_type_marker(φ))[ $(φ.val[1,1]), $(φ.val[1,2]), $(φ.val[1,3]);") + println(io, " $(φ.val[2,1]), $(φ.val[2,2]), $(φ.val[2,3]);") + print(io, " $(φ.val[3,1]), $(φ.val[3,2]), $(φ.val[3,3]) ]") end include("eucl/equi_coeffs_dict.jl") # Equivariant 3 x 3 Euclidean Matrix of general form (i.e., doesn't need to be symmetric or anti-symmetric) struct EuclideanMatrix{T} <: AbstractEuclideanMatrix{T} - val::SMatrix{3, 3, T, 9} + val::SMatrix{3,3,T,9} end _type_marker(φ::EuclideanMatrix) = "e" coco_init(phi::EuclideanMatrix{CT}, l, m, μ, T, A) where {CT<:Real} = ( - (l <= 2 && abs(m) <= l && abs(μ) <= l) - ? vec([EuclideanMatrix(conj.(mrmatrices[(l=l,m=-m,mu=-μ,i=i,j=j)])) for i=1:3 for j=1:3]) - : coco_zeros(phi, l, m, μ, T, A) ) + (l <= 2 && abs(m) <= l && abs(μ) <= l) + ? vec([EuclideanMatrix(conj.(mrmatrices[(l=l, m=-m, mu=-μ, i=i, j=j)])) for i = 1:3 for j = 1:3]) + : coco_zeros(phi, l, m, μ, T, A)) write_dict(φ::EuclideanMatrix{T}) where {T} = - Dict("__id__" => "ACEfrictionCore_EuclideanMatrix", - "valr" => write_dict(real.(Matrix(φ.val))), - "vali" => write_dict(imag.(Matrix(φ.val))), - "T" => write_dict(T) ) - -function read_dict(::Val{:ACEfrictionCore_EuclideanMatrix}, D::Dict) - T = read_dict(D["T"]) - valr = SMatrix{3, 3, T, 9}(read_dict(D["valr"])) - vali = SMatrix{3, 3, T, 9}(read_dict(D["vali"])) - return EuclideanMatrix{T}(valr + im * vali) + Dict("__id__" => "ACEfrictionCore_EuclideanMatrix", + "valr" => write_dict(real.(Matrix(φ.val))), + "vali" => write_dict(imag.(Matrix(φ.val))), + "T" => write_dict(T)) + +function read_dict(::Val{:ACEfrictionCore_EuclideanMatrix}, D::AbstractDict) + T = read_dict(D["T"]) + valr = SMatrix{3,3,T,9}(read_dict(D["valr"])) + vali = SMatrix{3,3,T,9}(read_dict(D["vali"])) + return EuclideanMatrix{T}(valr + im * vali) end # --------------------- -# Equivariant 3 x 3 symmetric Euclidean Matrix +# Equivariant 3 x 3 symmetric Euclidean Matrix -struct SymmetricEuclideanMatrix{T} <: AbstractEuclideanMatrix{T} #where {S<:MatrixSymmetry} - val::SMatrix{3, 3, T, 9} +struct SymmetricEuclideanMatrix{T} <: AbstractEuclideanMatrix{T} #where {S<:MatrixSymmetry} + val::SMatrix{3,3,T,9} end _type_marker(φ::SymmetricEuclideanMatrix) = "se" function coco_init(phi::SymmetricEuclideanMatrix{CT}, l, m, μ, T, A) where {CT<:Real} - return ( ( (l == 2 && abs(m) <= 2 && abs(μ) <= 2) || (l == 0 && abs(m) == 0 && abs(μ) == 0) ) - ? vec([SymmetricEuclideanMatrix(conj.(mrmatrices[(l=l,m=-m,mu=-μ,i=i,j=j)])) for i=1:3 for j=1:3]) - : coco_zeros(phi, l, m, μ, T, A) ) + return (((l == 2 && abs(m) <= 2 && abs(μ) <= 2) || (l == 0 && abs(m) == 0 && abs(μ) == 0)) + ? vec([SymmetricEuclideanMatrix(conj.(mrmatrices[(l=l, m=-m, mu=-μ, i=i, j=j)])) for i = 1:3 for j = 1:3]) + : coco_zeros(phi, l, m, μ, T, A)) end function ACEfrictionCore.write_dict(φ::SymmetricEuclideanMatrix{T}) where {T} - Dict("__id__" => "ACEfrictionCore_SymmetricEuclideanMatrix", - "valr" => write_dict(real.(Matrix(φ.val))), - "vali" => write_dict(imag.(Matrix(φ.val))), - "T" => write_dict(T)) -end - -function ACEfrictionCore.read_dict(::Val{:ACEfrictionCore_SymmetricEuclideanMatrix}, D::Dict) - T = read_dict(D["T"]) - valr = SMatrix{3, 3, T, 9}(read_dict(D["valr"])) - vali = SMatrix{3, 3, T, 9}(read_dict(D["vali"])) - return SymmetricEuclideanMatrix{T}(valr + im * vali) + Dict("__id__" => "ACEfrictionCore_SymmetricEuclideanMatrix", + "valr" => write_dict(real.(Matrix(φ.val))), + "vali" => write_dict(imag.(Matrix(φ.val))), + "T" => write_dict(T)) +end + +function ACEfrictionCore.read_dict(::Val{:ACEfrictionCore_SymmetricEuclideanMatrix}, D::AbstractDict) + T = read_dict(D["T"]) + valr = SMatrix{3,3,T,9}(read_dict(D["valr"])) + vali = SMatrix{3,3,T,9}(read_dict(D["vali"])) + return SymmetricEuclideanMatrix{T}(valr + im * vali) end -# Equivariant 3 x 3 anti-symmetric Euclidean Matrix +# Equivariant 3 x 3 anti-symmetric Euclidean Matrix # struct AntiSymmetricEuclideanMatrix{T} <: AbstractEuclideanMatrix{T} #where {S<:MatrixSymmetry} # val::SMatrix{3, 3, T, 9} # end @@ -370,10 +370,10 @@ end # Dict("__id__" => "ACEfrictionCore_AntiSymmetricEuclideanMatrix", # "valr" => write_dict(real.(Matrix(φ.val))), # "vali" => write_dict(imag.(Matrix(φ.val))), -# "T" => write_dict(T)) -# end +# "T" => write_dict(T)) +# end -# function ACEfrictionCore.read_dict(::Val{:ACEfrictionCore_AntiSymmetricEuclideanMatrix}, D::Dict) +# function ACEfrictionCore.read_dict(::Val{:ACEfrictionCore_AntiSymmetricEuclideanMatrix}, D::AbstractDict) # T = read_dict(D["T"]) # valr = SMatrix{3, 3, T, 9}(read_dict(D["valr"])) # vali = SMatrix{3, 3, T, 9}(read_dict(D["vali"])) @@ -381,53 +381,53 @@ end # end # for E = (:EuclideanMatrix,:SymmetricEuclideanMatrix, :AntiSymmetricEuclideanMatrix) -for E = (:EuclideanMatrix,:SymmetricEuclideanMatrix) +for E = (:EuclideanMatrix, :SymmetricEuclideanMatrix) eval( - quote - - complex(φ::$E) = $E(complex(φ.val)) - complex(::Type{$E{T}}) where {T} = $E{complex(T)} - - Base.convert(::Type{$E{T}}, φ::$E{T}) where {T<:Number} = $E(φ.val) #This function was previously not defined for EuclideanMatrix - isrealB(::$E{T}) where {T} = (T == real(T)) - - $E{T}() where {T <: Number} = $E{T}(zero(SMatrix{3, 3, T, 9})) - $E(T::DataType=Float64) = $E{T}() - coco_type(::Type{$E{T}}) where {T} = $E{complex(T)} - - Base.promote_rule(::Type{T1}, ::Type{$E{T2}} - ) where {T1 <: Number, T2 <: Number} = - $E{promote_type(T1, T2)} - - end - ) - end - - - Base.promote_rule(t1::Type{T1}, t2::Type{SMatrix{N,N, T2}} - ) where {N, T1 <: Number, T2 <: AbstractProperty} = - SMatrix{N, N, promote_rule(t1, T2)} - -# --------------------- -# some codes to help convert from svector, smatrix to spherical ... - -@generated function __getL(::Val{N}) where {N} - @assert isodd(N) - quote - $( div(N-1, 2) ) - end + quote + + complex(φ::$E) = $E(complex(φ.val)) + complex(::Type{$E{T}}) where {T} = $E{complex(T)} + + Base.convert(::Type{$E{T}}, φ::$E{T}) where {T<:Number} = $E(φ.val) #This function was previously not defined for EuclideanMatrix + isrealB(::$E{T}) where {T} = (T == real(T)) + + $E{T}() where {T<:Number} = $E{T}(zero(SMatrix{3,3,T,9})) + $E(T::DataType=Float64) = $E{T}() + coco_type(::Type{$E{T}}) where {T} = $E{complex(T)} + + Base.promote_rule(::Type{T1}, ::Type{$E{T2}} + ) where {T1<:Number,T2<:Number} = + $E{promote_type(T1, T2)} + + end + ) +end + + +Base.promote_rule(t1::Type{T1}, t2::Type{SMatrix{N,N,T2}} +) where {N,T1<:Number,T2<:AbstractProperty} = + SMatrix{N,N,promote_rule(t1, T2)} + +# --------------------- +# some codes to help convert from svector, smatrix to spherical ... + +@generated function __getL(::Val{N}) where {N} + @assert isodd(N) + quote + $(div(N - 1, 2)) + end end __getL(::SVector{N}) where {N} = __getL(Val{N}()) -__getL1L2(::SMatrix{N1, N2}) where {N1, N2} = __getL(Val{N1}()), __getL(Val{N2}()) +__getL1L2(::SMatrix{N1,N2}) where {N1,N2} = __getL(Val{N1}()), __getL(Val{N2}()) # --------------------- SphericalVector # TODO: remove the Val{L} - not needed (next cleanup) -struct SphericalVector{L, LEN, T} <: AbstractProperty - val::SVector{LEN, T} - _valL::Val{L} +struct SphericalVector{L,LEN,T} <: AbstractProperty + val::SVector{LEN,T} + _valL::Val{L} end SphericalVector(val::SVector) = SphericalVector(val, Val(__getL(val))) @@ -436,8 +436,8 @@ SphericalVector(val::SVector) = SphericalVector(val, Val(__getL(val))) # # differentiation - cf #27 # *(φ::SphericalVector, dAA::SVector) = φ.val * dAA' -isrealB(::SphericalVector) = false -isrealAA(::SphericalVector) = false +isrealB(::SphericalVector) = false +isrealAA(::SphericalVector) = false real(φ::SphericalVector) = SphericalVector(real(φ.val), φ._valL) @@ -449,29 +449,29 @@ getL(φ::SphericalVector{L}) where {L} = L # L = 3 -> ... + 5 -> 9 # 1 + 3 + 5 + ... + 2*L+1 # = L + 2 * (1 + ... + L) = L+1 + 2 * L * (L+1) / 2 = (L+1)^2 -function SphericalVector(L::Integer; T = Float64) - LEN = 2L+1 # length of SH basis up to L - return SphericalVector( zero(SVector{LEN, T}), Val{L}() ) +function SphericalVector(L::Integer; T=Float64) + LEN = 2L + 1 # length of SH basis up to L + return SphericalVector(zero(SVector{LEN,T}), Val{L}()) end -function SphericalVector{L, LEN, T}(x::AbstractArray) where {L, LEN, T} - @assert length(x) == LEN - SphericalVector{L, LEN, T}( SVector{LEN, T}(x), Val(L) ) +function SphericalVector{L,LEN,T}(x::AbstractArray) where {L,LEN,T} + @assert length(x) == LEN + SphericalVector{L,LEN,T}(SVector{LEN,T}(x), Val(L)) end -SphericalVector{L, LEN, T}() where {L, LEN, T} = - SphericalVector( zero(SVector{LEN, T}), Val{L}() ) +SphericalVector{L,LEN,T}() where {L,LEN,T} = + SphericalVector(zero(SVector{LEN,T}), Val{L}()) function filter(φ::SphericalVector, grp::O3, b::Array) - if length(b) <= 1 - return true - end - suml = sum( getl(grp, bi) for bi in b ) - if haskey(b[1], msym(grp)) - summ = sum( getm(grp, bi) for bi in b ) - return iseven(suml) == iseven(getL(φ)) && abs(summ) <= getL(φ) - end - return iseven(suml) == iseven(getL(φ)) + if length(b) <= 1 + return true + end + suml = sum(getl(grp, bi) for bi in b) + if haskey(b[1], msym(grp)) + summ = sum(getm(grp, bi) for bi in b) + return iseven(suml) == iseven(getL(φ)) && abs(summ) <= getL(φ) + end + return iseven(suml) == iseven(getL(φ)) end rot3Dcoeffs(::SphericalVector, T::DataType=Float64) = Rot3DCoeffs(T) @@ -484,66 +484,68 @@ using ACEfrictionCore.Wigner: wigner_D_indices # Equation (1.2) - vector value coupling coefficients # ∫_{SO3} D^{ll}_{μμmm} D^*(Q) e^t dQ -> 2L+1 column vector function vec_cou_coe(rotc::Rot3DCoeffs{T}, - l::Integer, m::Integer, μ::Integer, - L::Integer, t::Integer) where {T} - @assert 0 < t <= 2L+1 - D = wigner_D_indices(L)' # Dt = D[:,t] --> # D^* ⋅ e^t - LL = SA[l, L] - Z = ntuple(i -> begin - cc = (rotc(LL, SA[μ, D[i, t].m], SA[m, D[i, t].μ]).val)::T - D[i, t].sign * cc - end, 2*L+1) - return SphericalVector{L, 2L+1, Complex{T}}(SVector(Z)) + l::Integer, m::Integer, μ::Integer, + L::Integer, t::Integer) where {T} + @assert 0 < t <= 2L + 1 + D = wigner_D_indices(L)' # Dt = D[:,t] --> # D^* ⋅ e^t + LL = SA[l, L] + Z = ntuple(i -> begin + cc = (rotc(LL, SA[μ, D[i, t].m], SA[m, D[i, t].μ]).val)::T + D[i, t].sign * cc + end, 2 * L + 1) + return SphericalVector{L,2L + 1,Complex{T}}(SVector(Z)) end function _select_t(φ::SphericalVector{L}, l, M, K) where {L} - D = wigner_D_indices(L)' - tret = -1; numt = 0 - for t = 1:2L+1 - prodμt = prod( (D[i, t].μ + M) for i in 1:2L+1) # avoid more allocations - prodmt = prod( (D[i, t].m + K) for i in 1:2L+1) - if prodμt == prodmt == 0 - tret = t; numt += 1 - end - end - # We assumed that there is only one coefficient; this will warn us if it fails - @assert numt == 1 - return tret + D = wigner_D_indices(L)' + tret = -1 + numt = 0 + for t = 1:2L+1 + prodμt = prod((D[i, t].μ + M) for i in 1:2L+1) # avoid more allocations + prodmt = prod((D[i, t].m + K) for i in 1:2L+1) + if prodμt == prodmt == 0 + tret = t + numt += 1 + end + end + # We assumed that there is only one coefficient; this will warn us if it fails + @assert numt == 1 + return tret end -coco_zeros(φ::TP, ll, mm, kk, T, A) where {TP <: SphericalVector} = zero(TP) +coco_zeros(φ::TP, ll, mm, kk, T, A) where {TP<:SphericalVector} = zero(TP) coco_dot(u1::SphericalVector, u2::SphericalVector) = - dot(u1.val, u2.val) + dot(u1.val, u2.val) coco_filter(φ::SphericalVector{L}, ll, mm) where {L} = - iseven(sum(ll) + L) && (abs(sum(mm)) <= L) + iseven(sum(ll) + L) && (abs(sum(mm)) <= L) coco_filter(φ::SphericalVector{L}, ll, mm, kk) where {L} = - iseven(sum(ll) + L) && (abs(sum(mm)) <= L) && (abs(sum(kk)) <= L) + iseven(sum(ll) + L) && (abs(sum(mm)) <= L) && (abs(sum(kk)) <= L) coco_init(φ::SphericalVector{L}, l, m, μ, T, A) where {L} = - vec_cou_coe(__rotcoeff_inv, l, m, μ, L, _select_t(φ, l, m, μ)) + vec_cou_coe(__rotcoeff_inv, l, m, μ, L, _select_t(φ, l, m, μ)) -coco_init(φ::SphericalVector{L}) where {L} = - L==0 ? reshape([SphericalVector(SVector(1.0+0.0im))],1,1) : [] +coco_init(φ::SphericalVector{L}) where {L} = + L == 0 ? reshape([SphericalVector(SVector(1.0 + 0.0im))], 1, 1) : [] # --------------- SphericalMatrix -struct SphericalMatrix{L1, L2, LEN1, LEN2, T, LL} <: AbstractProperty - val::SMatrix{LEN1, LEN2, T, LL} - _valL1::Val{L1} - _valL2::Val{L2} +struct SphericalMatrix{L1,L2,LEN1,LEN2,T,LL} <: AbstractProperty + val::SMatrix{LEN1,LEN2,T,LL} + _valL1::Val{L1} + _valL2::Val{L2} end -function Base.show(io::IO, φ::Union{SphericalVector, SphericalMatrix}) - buffer = IOBuffer() - print(buffer, round.(φ.val, digits=4)) - str = String(take!(buffer)) - lenn = findfirst('[', str) - print(io, "y" * str[lenn:end]) +function Base.show(io::IO, φ::Union{SphericalVector,SphericalMatrix}) + buffer = IOBuffer() + print(buffer, round.(φ.val, digits=4)) + str = String(take!(buffer)) + lenn = findfirst('[', str) + print(io, "y" * str[lenn:end]) end SphericalMatrix(val::SMatrix) = SphericalMatrix(val, Val.(__getL1L2(val))...) @@ -557,8 +559,8 @@ SphericalMatrix(val::SMatrix) = SphericalMatrix(val, Val.(__getL1L2(val))...) getL(φ::SphericalMatrix{L1,L2}) where {L1,L2} = L1, L2 -isrealB(::SphericalMatrix) = false -isrealAA(::SphericalMatrix) = false +isrealB(::SphericalMatrix) = false +isrealAA(::SphericalMatrix) = false # L = 0 -> (0,0) @@ -566,161 +568,158 @@ isrealAA(::SphericalMatrix) = false # L = 3 -> ... + 5 -> 9 # 1 + 3 + 5 + ... + 2*L+1 # = L + 2 * (1 + ... + L) = L+1 + 2 * L * (L+1) / 2 = (L+1)^2 -function SphericalMatrix(L1::Integer, L2::Integer; T = Float64) - LEN1 = 2L1+1 # length of SH basis up to L - LEN2 = 2L2+1 - return SphericalMatrix( zero(SMatrix{LEN1, LEN2, T}), Val{L1}(), Val{L2}() ) +function SphericalMatrix(L1::Integer, L2::Integer; T=Float64) + LEN1 = 2L1 + 1 # length of SH basis up to L + LEN2 = 2L2 + 1 + return SphericalMatrix(zero(SMatrix{LEN1,LEN2,T}), Val{L1}(), Val{L2}()) end # this is a mess - need to fix it, we should need just one constructor??? -function SphericalMatrix{L1, L2, LEN1, LEN2, T, LL}(x::AbstractMatrix) where {L1, L2, LEN1, LEN2, T, LL} - @assert size(x) == (LEN1, LEN2) - SphericalMatrix{L1, L2, LEN1, LEN2, T, LL}( SMatrix{LEN1, LEN2, T}(x), Val(L1), Val(L2) ) +function SphericalMatrix{L1,L2,LEN1,LEN2,T,LL}(x::AbstractMatrix) where {L1,L2,LEN1,LEN2,T,LL} + @assert size(x) == (LEN1, LEN2) + SphericalMatrix{L1,L2,LEN1,LEN2,T,LL}(SMatrix{LEN1,LEN2,T}(x), Val(L1), Val(L2)) end -SphericalMatrix{L1, L2, LEN1, LEN2, T, LL}() where {L1, L2, LEN1, LEN2, T, LL} = - SphericalMatrix( zero(SMatrix{LEN1, LEN2, T}), Val{L1}(), Val{L2}() ) +SphericalMatrix{L1,L2,LEN1,LEN2,T,LL}() where {L1,L2,LEN1,LEN2,T,LL} = + SphericalMatrix(zero(SMatrix{LEN1,LEN2,T}), Val{L1}(), Val{L2}()) -SphericalMatrix{L1, L2, LEN1, LEN2, T}() where {L1, L2, LEN1, LEN2, T} = - SphericalMatrix( zero(SMatrix{LEN1, LEN2, T}), Val{L1}(), Val{L2}() ) +SphericalMatrix{L1,L2,LEN1,LEN2,T}() where {L1,L2,LEN1,LEN2,T} = + SphericalMatrix(zero(SMatrix{LEN1,LEN2,T}), Val{L1}(), Val{L2}()) function filter(φ::SphericalMatrix, grp::O3, b::Array) - if length(b) < 1 - return true - end - suml = sum( getl(grp, bi) for bi in b ) - if haskey(b[1], msym(grp)) - summ = sum( getm(grp, bi) for bi in b ) - return iseven(suml) == iseven( sum(getL(φ)) ) && abs(summ) <= sum(getL(φ)) - end - return iseven(suml) == iseven( sum(getL(φ)) ) + if length(b) < 1 + return true + end + suml = sum(getl(grp, bi) for bi in b) + if haskey(b[1], msym(grp)) + summ = sum(getm(grp, bi) for bi in b) + return iseven(suml) == iseven(sum(getL(φ))) && abs(summ) <= sum(getL(φ)) + end + return iseven(suml) == iseven(sum(getL(φ))) end rot3Dcoeffs(::SphericalMatrix, T::DataType=Float64) = Rot3DCoeffs(T) function mat_cou_coe(rotc::Rot3DCoeffs{T}, - l::Integer, m::Integer, μ::Integer, - a::Integer, b::Integer, - ::Val{L1}, ::Val{L2}) where {T, L1, L2} - @assert (0 < a <= 2L1 + 1) && (0 < b <= 2L2 + 1) - Z = zero(MMatrix{2L1+1, 2L2+1, Complex{T}}) # zeros(2 * L1 + 1, 2 * L2 + 1) - Dp = wigner_D_indices(L1)' - Dq = wigner_D_indices(L2) - LL = SA[l, L1, L2] - for i = 1:(2 * L1 + 1) - for j = 1:(2 * L2 + 1) - MM = SA[μ, Dp[i,a].m, Dq[b,j].m] - KK = SA[m, Dp[i,a].μ, Dq[b,j].μ] - cc = (rotc(LL, MM, KK).val)::T - Z[i,j] = Dp[i,a].sign * Dq[b,j].sign * cc - end - end - return SphericalMatrix(SMatrix(Z), Val{L1}(), Val{L2}()) - # return SphericalMatrix(SMatrix{2L1+1,2L2+1,Complex{T}}(Z), Val{L1}(), Val{L2}()) + l::Integer, m::Integer, μ::Integer, + a::Integer, b::Integer, + ::Val{L1}, ::Val{L2}) where {T,L1,L2} + @assert (0 < a <= 2L1 + 1) && (0 < b <= 2L2 + 1) + Z = zero(MMatrix{2L1 + 1,2L2 + 1,Complex{T}}) # zeros(2 * L1 + 1, 2 * L2 + 1) + Dp = wigner_D_indices(L1)' + Dq = wigner_D_indices(L2) + LL = SA[l, L1, L2] + for i = 1:(2*L1+1) + for j = 1:(2*L2+1) + MM = SA[μ, Dp[i, a].m, Dq[b, j].m] + KK = SA[m, Dp[i, a].μ, Dq[b, j].μ] + cc = (rotc(LL, MM, KK).val)::T + Z[i, j] = Dp[i, a].sign * Dq[b, j].sign * cc + end + end + return SphericalMatrix(SMatrix(Z), Val{L1}(), Val{L2}()) + # return SphericalMatrix(SMatrix{2L1+1,2L2+1,Complex{T}}(Z), Val{L1}(), Val{L2}()) end function _select_ab(φ::SphericalMatrix{L1,L2}, M, K) where {L1,L2} - Dp = wigner_D_indices(L1)' - Dq = wigner_D_indices(L2) - list_ab = Tuple{Int, Int}[] - for a = 1:2L1+1 - for b = 1:2L2+1 - # pm = prod( ma[i] + mb[j] + K for i = 1:2L1+1, j = 1:2L2+1) - pm = prod( Dp[i,a].m + Dq[b,j].m + K for i = 1:2L1+1, j = 1:2L2+1) - # pμ = prod(μa[i] + μb[j] + M for i = 1:2L1+1, j = 1:2L2+1) - pμ = prod( Dp[i,a].μ + Dq[b,j].μ + M for i = 1:2L1+1, j = 1:2L2+1) - if pμ == pm ==0 - push!(list_ab, (a,b)) - end - end - end - return list_ab + Dp = wigner_D_indices(L1)' + Dq = wigner_D_indices(L2) + list_ab = Tuple{Int,Int}[] + for a = 1:2L1+1 + for b = 1:2L2+1 + # pm = prod( ma[i] + mb[j] + K for i = 1:2L1+1, j = 1:2L2+1) + pm = prod(Dp[i, a].m + Dq[b, j].m + K for i = 1:2L1+1, j = 1:2L2+1) + # pμ = prod(μa[i] + μb[j] + M for i = 1:2L1+1, j = 1:2L2+1) + pμ = prod(Dp[i, a].μ + Dq[b, j].μ + M for i = 1:2L1+1, j = 1:2L2+1) + if pμ == pm == 0 + push!(list_ab, (a, b)) + end + end + end + return list_ab end -function coco_init(φ::SphericalMatrix{L1,L2}, l, m, μ, T, A) where{L1,L2} - list = _select_ab(φ, m, μ) - @assert length(list) > 0 - # MAIN CASE - if abs(m) <= L1+L2 && abs(μ) <= L1+L2 - return [ mat_cou_coe(__rotcoeff_inv, l, m, μ, a, b, Val(L1), Val(L2)) - for (a,b) in list ] - end +function coco_init(φ::SphericalMatrix{L1,L2}, l, m, μ, T, A) where {L1,L2} + list = _select_ab(φ, m, μ) + @assert length(list) > 0 + # MAIN CASE + if abs(m) <= L1 + L2 && abs(μ) <= L1 + L2 + return [mat_cou_coe(__rotcoeff_inv, l, m, μ, a, b, Val(L1), Val(L2)) + for (a, b) in list] + end - # @warn("SHOULDN'T BE HERE!!") - return fill( zero(typeof(φ)), length(list) ) + # @warn("SHOULDN'T BE HERE!!") + return fill(zero(typeof(φ)), length(list)) end -coco_init(φ::SphericalMatrix{L1,L2}) where{L1,L2} = - L1==L2 ? reshape([ACEfrictionCore.SphericalMatrix(SMatrix{2L1+1,2L2+1,ComplexF64}(I(2L1+1)),Val(L1),Val(L2))],1,1) : [] +coco_init(φ::SphericalMatrix{L1,L2}) where {L1,L2} = + L1 == L2 ? reshape([ACEfrictionCore.SphericalMatrix(SMatrix{2L1 + 1,2L2 + 1,ComplexF64}(I(2L1 + 1)), Val(L1), Val(L2))], 1, 1) : [] -coco_zeros(φ::TP, ll, mm, kk, T, A) where{TP <: SphericalMatrix} = - zeros(TP, length(_select_ab(φ, sum(mm), sum(kk)))) +coco_zeros(φ::TP, ll, mm, kk, T, A) where {TP<:SphericalMatrix} = + zeros(TP, length(_select_ab(φ, sum(mm), sum(kk)))) coco_filter(φ::SphericalMatrix{L1,L2}, ll, mm) where {L1,L2} = - iseven(sum(ll)) == iseven(L1+L2) && (abs(sum(mm)) <= L1+L2) + iseven(sum(ll)) == iseven(L1 + L2) && (abs(sum(mm)) <= L1 + L2) coco_filter(φ::SphericalMatrix{L1,L2}, ll, mm, kk) where {L1,L2} = - iseven(sum(ll) + L1 + L2) && - (abs(sum(mm)) <= L1+L2) && - (abs(sum(kk)) <= L1+L2) + iseven(sum(ll) + L1 + L2) && + (abs(sum(mm)) <= L1 + L2) && + (abs(sum(kk)) <= L1 + L2) coco_dot(u1::SphericalMatrix, u2::SphericalMatrix) = - dot(u1.val, u2.val) + dot(u1.val, u2.val) -# Some promotion rules +# Some promotion rules Base.promote_rule(::Type{T1}, ::Type{Invariant{T2}} - ) where {T1 <: Number, T2 <: Number} = - Invariant{promote_type(T1, T2)} +) where {T1<:Number,T2<:Number} = + Invariant{promote_type(T1, T2)} Base.promote_rule(::Type{T1}, ::Type{EuclideanVector{T2}} - ) where {T1 <: Number, T2 <: Number} = - EuclideanVector{promote_type(T1, T2)} +) where {T1<:Number,T2<:Number} = + EuclideanVector{promote_type(T1, T2)} -Base.promote_rule(t1::Type{T1}, t2::Type{SVector{N, T2}} - ) where {N, T1 <: Number, T2 <: AbstractProperty} = - SVector{N, promote_rule(t1, T2)} +Base.promote_rule(t1::Type{T1}, t2::Type{SVector{N,T2}} +) where {N,T1<:Number,T2<:AbstractProperty} = + SVector{N,promote_rule(t1, T2)} -# --------------------------- AD related codes +# --------------------------- AD related codes -# an x -> x.val implementation with custom adjoints to sort out the +# an x -> x.val implementation with custom adjoints to sort out the # mess created by the AbstractProperties -# maybe this feels a bit wrong, definitely a hack. What might be nicer -# is to introduce a "Dual Property" similar to the "DState"; Then we -# could have something along the lines of DProp * Prop = scalar or -# _contract(DProp, Prop) = scalar; That would be the "systematic" and -# "disciplined" way of implementing this. +# maybe this feels a bit wrong, definitely a hack. What might be nicer +# is to introduce a "Dual Property" similar to the "DState"; Then we +# could have something along the lines of DProp * Prop = scalar or +# _contract(DProp, Prop) = scalar; That would be the "systematic" and +# "disciplined" way of implementing this. """ -`val(x) = x.val`, normally to be used if x is a property. This should be used -instead of x.val when the operation is part of a bigger expression that is -to be ADed. I.e. `val` has rrules implemented that should allow taking up -to two derivatives. +`val(x) = x.val`, normally to be used if x is a property. This should be used +instead of x.val when the operation is part of a bigger expression that is +to be ADed. I.e. `val` has rrules implemented that should allow taking up +to two derivatives. TODO: at the moment this is a bit hacky, and needs to be adjusted over time -as we learn more about how to best implement AD. The real question is what -is the correct adjoint of this operation? E.g., if +as we learn more about how to best implement AD. The real question is what +is the correct adjoint of this operation? E.g., if ```` val : { Invariant{T} } -> { T } ``` -then should +then should ``` val* : { T } -> { Invariant{T} } ``` -? If this is the case, then we need a parameterised `val` function. I.e. +? If this is the case, then we need a parameterised `val` function. I.e. we need to remember what the Property is that we started from. """ -val(x) = x.val +val(x) = x.val -# _rrule_val and rrule functions removed - derivative functionality has been removed +# _rrule_val and rrule functions removed - derivative functionality has been removed # import ChainRulesCore: ProjectTo # (::ProjectTo{T})(φ::Invariant{T}) where {T} = val(φ) - - - diff --git a/src/rotations3d.jl b/src/rotations3d.jl index 92c0223..25e1420 100644 --- a/src/rotations3d.jl +++ b/src/rotations3d.jl @@ -5,8 +5,8 @@ module Rotations3D using StaticArrays using LinearAlgebra: norm, rank, svd, Diagonal, tr -using ACEfrictionCore: coco_zeros, coco_init, coco_dot, coco_filter, AbstractProperty, - coco_type +using ACEfrictionCore: coco_zeros, coco_init, coco_dot, coco_filter, AbstractProperty, + coco_type export ClebschGordan, Rot3DCoeffs, ri_basis, rpi_basis, R3DC, Rot3DCoeffsEquiv @@ -15,54 +15,54 @@ export ClebschGordan, Rot3DCoeffs, ri_basis, rpi_basis, R3DC, Rot3DCoeffsEquiv `?clebschgordan` for the convention that is use. """ struct ClebschGordan{T} - vals::Dict{Tuple{Int, Int, Int, Int, Int, Int}, T} + vals::AbstractDict{Tuple{Int,Int,Int,Int,Int,Int},T} end # -> CouplingCoeffRecursion ???? -struct Rot3DCoeffs{T, TP} - vals::Vector{Dict} # val[N] = coeffs for correlation order N - cg::ClebschGordan{T} - phi::TP +struct Rot3DCoeffs{T,TP} + vals::Vector{Dict} # val[N] = coeffs for correlation order N + cg::ClebschGordan{T} + phi::TP end # ----------------------------------- # iterating over an m collection # ----------------------------------- -_mvec(::CartesianIndex{0}) = SVector{0, Int}() +_mvec(::CartesianIndex{0}) = SVector{0,Int}() _mvec(mpre::CartesianIndex) = SVector(Tuple(mpre)...) -struct MRange{N, T2, TP} - ll::SVector{N, Int} - cartrg::T2 - phi::TP +struct MRange{N,T2,TP} + ll::SVector{N,Int} + cartrg::T2 + phi::TP end Base.length(mr::MRange) = - sum(mt -> coco_filter(mr.phi, mr.ll, _mvec(mt)), mr.cartrg) + sum(mt -> coco_filter(mr.phi, mr.ll, _mvec(mt)), mr.cartrg) """ Given an l-vector `ll` iterate over all combinations of `mm` vectors of the same length such that `sum(mm) == 0` """ _mrange(phi, ll) = - MRange(ll, CartesianIndices(ntuple(i -> -ll[i]:ll[i], length(ll))), phi) + MRange(ll, CartesianIndices(ntuple(i -> -ll[i]:ll[i], length(ll))), phi) # TODO: should we impose here that (ll, mm) are lexicographically ordered? function Base.iterate(mr::MRange, idx::Integer=0) - while true - idx += 1 - if idx > length(mr.cartrg) - return nothing - end - mm = _mvec(mr.cartrg[idx]) - if coco_filter(mr.phi, mr.ll, mm) - return mm, idx - end - end - error("we should never be here") + while true + idx += 1 + if idx > length(mr.cartrg) + return nothing + end + mm = _mvec(mr.cartrg[idx]) + if coco_filter(mr.phi, mr.ll, mm) + return mm, idx + end + end + error("we should never be here") end @@ -73,12 +73,12 @@ end # ---------------------------------------------------------------------- -cg_conditions(j1,m1, j2,m2, J,M) = - cg_l_condition(j1, j2, J) && - cg_m_condition(m1, m2, M) && - (abs(m1) <= j1) && (abs(m2) <= j2) && (abs(M) <= J) +cg_conditions(j1, m1, j2, m2, J, M) = + cg_l_condition(j1, j2, J) && + cg_m_condition(m1, m2, M) && + (abs(m1) <= j1) && (abs(m2) <= j2) && (abs(M) <= J) -cg_l_condition(j1, j2, J) = (abs(j1-j2) <= J <= j1 + j2) +cg_l_condition(j1, j2, J) = (abs(j1 - j2) <= J <= j1 + j2) cg_m_condition(m1, m2, M) = (M == m1 + m2) @@ -106,53 +106,53 @@ where ``` """ function clebschgordan(j1, m1, j2, m2, J, M, T=Float64) - if !cg_conditions(j1, m1, j2, m2, J, M) - return zero(T) - end - - N = (2*J+1) * - factorial(big(j1+m1)) * factorial(big(j1-m1)) * - factorial(big(j2+m2)) * factorial(big(j2-m2)) * - factorial(big(J+M)) * factorial(big(J-M)) / - factorial(big( j1+j2-J)) / - factorial(big( j1-j2+J)) / - factorial(big(-j1+j2+J)) / - factorial(big(j1+j2+J+1)) - - G = big(0) - # 0 ≦ k ≦ j1+j2-J - # 0 ≤ j1-m1-k ≤ j1-j2+J <=> j2-J-m1 ≤ k ≤ j1-m1 - # 0 ≤ j2+m2-k ≤ -j1+j2+J <=> j1-J+m2 ≤ k ≤ j2+m2 - lb = (0, j2-J-m1, j1-J+m2) - ub = (j1+j2-J, j1-m1, j2+m2) - for k in maximum(lb):minimum(ub) - bk = big(k) - G += (-1)^k * - binomial(big( j1+j2-J), big(k)) * - binomial(big( j1-j2+J), big(j1-m1-k)) * - binomial(big(-j1+j2+J), big(j2+m2-k)) - end - - return T(sqrt(N) * G) + if !cg_conditions(j1, m1, j2, m2, J, M) + return zero(T) + end + + N = (2 * J + 1) * + factorial(big(j1 + m1)) * factorial(big(j1 - m1)) * + factorial(big(j2 + m2)) * factorial(big(j2 - m2)) * + factorial(big(J + M)) * factorial(big(J - M)) / + factorial(big(j1 + j2 - J)) / + factorial(big(j1 - j2 + J)) / + factorial(big(-j1 + j2 + J)) / + factorial(big(j1 + j2 + J + 1)) + + G = big(0) + # 0 ≦ k ≦ j1+j2-J + # 0 ≤ j1-m1-k ≤ j1-j2+J <=> j2-J-m1 ≤ k ≤ j1-m1 + # 0 ≤ j2+m2-k ≤ -j1+j2+J <=> j1-J+m2 ≤ k ≤ j2+m2 + lb = (0, j2 - J - m1, j1 - J + m2) + ub = (j1 + j2 - J, j1 - m1, j2 + m2) + for k in maximum(lb):minimum(ub) + bk = big(k) + G += (-1)^k * + binomial(big(j1 + j2 - J), big(k)) * + binomial(big(j1 - j2 + J), big(j1 - m1 - k)) * + binomial(big(-j1 + j2 + J), big(j2 + m2 - k)) + end + + return T(sqrt(N) * G) end ClebschGordan(T=Float64) = - ClebschGordan{T}(Dict{Tuple{Int,Int,Int,Int,Int,Int}, T}()) + ClebschGordan{T}(Dict{Tuple{Int,Int,Int,Int,Int,Int},T}()) _cg_key(j1, m1, j2, m2, J, M) = (j1, m1, j2, m2, J, M) function (cg::ClebschGordan{T})(j1, m1, j2, m2, J, M) where {T} - if !cg_conditions(j1,m1, j2,m2, J,M) - return zero(T) - end - key = _cg_key(j1, m1, j2, m2, J, M) - if haskey(cg.vals, key) - return cg.vals[key] - end - val = clebschgordan(j1, m1, j2, m2, J, M, T) - cg.vals[key] = val - return val + if !cg_conditions(j1, m1, j2, m2, J, M) + return zero(T) + end + key = _cg_key(j1, m1, j2, m2, J, M) + if haskey(cg.vals, key) + return cg.vals[key] + end + val = clebschgordan(j1, m1, j2, m2, J, M, T) + cg.vals[key] = val + return val end @@ -166,40 +166,40 @@ end dicttype(N::Integer, TP) = dicttype(Val(N), TP) dicttype(::Val{N}, TP) where {N} = - Dict{Tuple{SVector{N,Int}, SVector{N,Int}, SVector{N,Int}}, TP} + Dict{Tuple{SVector{N,Int},SVector{N,Int},SVector{N,Int}},TP} Rot3DCoeffs(φ, T=Float64) = Rot3DCoeffs(Dict[], ClebschGordan(T), φ) function get_vals(A::Rot3DCoeffs{T}, valN::Val{N}) where {T,N} - # make up an ll, kk, mm and compute a dummy coupling coeff - ll, mm, kk = SVector(0), SVector(0), SVector(0) - cc0 = coco_zeros(A.phi, ll, mm, kk, T, A) - TP = typeof(cc0) - if length(A.vals) < N - # create more dictionaries of the correct type - for n = length(A.vals)+1:N - push!(A.vals, dicttype(n, TP)()) - end - end - return (A.vals[N])::(dicttype(valN, TP)) + # make up an ll, kk, mm and compute a dummy coupling coeff + ll, mm, kk = SVector(0), SVector(0), SVector(0) + cc0 = coco_zeros(A.phi, ll, mm, kk, T, A) + TP = typeof(cc0) + if length(A.vals) < N + # create more dictionaries of the correct type + for n = length(A.vals)+1:N + push!(A.vals, dicttype(n, TP)()) + end + end + return (A.vals[N])::(dicttype(valN, TP)) end _key(ll::StaticVector{N}, mm::StaticVector{N}, kk::StaticVector{N}) where {N} = - (SVector{N, Int}(ll), SVector{N, Int}(mm), SVector{N, Int}(kk)) + (SVector{N,Int}(ll), SVector{N,Int}(mm), SVector{N,Int}(kk)) function (A::Rot3DCoeffs{T})(ll::StaticVector{N}, - mm::StaticVector{N}, - kk::StaticVector{N}) where {T, N} - vals = get_vals(A, Val(N)) # this should infer the type! - key = _key(ll, mm, kk) - if haskey(vals, key) - val = vals[key] - else - val = _compute_val(A, key...) - vals[key] = val - end - return val + mm::StaticVector{N}, + kk::StaticVector{N}) where {T,N} + vals = get_vals(A, Val(N)) # this should infer the type! + key = _key(ll, mm, kk) + if haskey(vals, key) + val = vals[key] + else + val = _compute_val(A, key...) + vals[key] = val + end + return val end # the recursion has two steps so we need to define the @@ -208,43 +208,43 @@ end # or reshuffling should allow us to get rid of the {N = 2} case. (A::Rot3DCoeffs{T})(ll::StaticVector{1}, - mm::StaticVector{1}, - kk::StaticVector{1}) where {T} = - coco_init(A.phi, ll[1], mm[1], kk[1], T, A) + mm::StaticVector{1}, + kk::StaticVector{1}) where {T} = + coco_init(A.phi, ll[1], mm[1], kk[1], T, A) function _compute_val(A::Rot3DCoeffs{T}, ll::StaticVector{N}, - mm::StaticVector{N}, - kk::StaticVector{N}) where {T, N} - val = coco_zeros(A.phi, ll, mm, kk, T, A) - TV = typeof(val) - - tmp = zero(MVector{N-1, Int}) - - function _get_pp(aa, ap) - for i = 1:N-2 - @inbounds tmp[i] = aa[i] - end - tmp[N-1] = ap - return SVector(tmp) - end - - jmin = maximum( ( abs(ll[N-1]-ll[N]), - abs(kk[N-1]+kk[N]), - abs(mm[N-1]+mm[N]) ) ) - jmax = ll[N-1]+ll[N] - for j = jmin:jmax - cgk = A.cg(ll[N-1], kk[N-1], ll[N], kk[N], j, kk[N-1]+kk[N]) - cgm = A.cg(ll[N-1], mm[N-1], ll[N], mm[N], j, mm[N-1]+mm[N]) - if cgk * cgm != 0 - llpp = _get_pp(ll, j) # SVector(llp..., j) - mmpp = _get_pp(mm, mm[N-1]+mm[N]) # SVector(mmp..., mm[N-1]+mm[N]) - kkpp = _get_pp(kk, kk[N-1]+kk[N]) # SVector(kkp..., kk[N-1]+kk[N]) - a = A(llpp, mmpp, kkpp)::TV - val += cgk * cgm * a - end - end - return val + mm::StaticVector{N}, + kk::StaticVector{N}) where {T,N} + val = coco_zeros(A.phi, ll, mm, kk, T, A) + TV = typeof(val) + + tmp = zero(MVector{N - 1,Int}) + + function _get_pp(aa, ap) + for i = 1:N-2 + @inbounds tmp[i] = aa[i] + end + tmp[N-1] = ap + return SVector(tmp) + end + + jmin = maximum((abs(ll[N-1] - ll[N]), + abs(kk[N-1] + kk[N]), + abs(mm[N-1] + mm[N]))) + jmax = ll[N-1] + ll[N] + for j = jmin:jmax + cgk = A.cg(ll[N-1], kk[N-1], ll[N], kk[N], j, kk[N-1] + kk[N]) + cgm = A.cg(ll[N-1], mm[N-1], ll[N], mm[N], j, mm[N-1] + mm[N]) + if cgk * cgm != 0 + llpp = _get_pp(ll, j) # SVector(llp..., j) + mmpp = _get_pp(mm, mm[N-1] + mm[N]) # SVector(mmp..., mm[N-1]+mm[N]) + kkpp = _get_pp(kk, kk[N-1] + kk[N]) # SVector(kkp..., kk[N-1]+kk[N]) + a = A(llpp, mmpp, kkpp)::TV + val += cgk * cgm * a + end + end + return val end # ---------------------------------------------------------------------- @@ -255,77 +255,77 @@ end function re_basis(A::Rot3DCoeffs{T}, ll::SVector) where {T} - TP = typeof(A.phi) - TCC = coco_type(TP) - CC, Mll = compute_Al(A, ll) # CC::Vector{Vector{...}} - G = [ sum( coco_dot(CC[a][i], CC[b][i]) for i = 1:length(Mll) ) - for a = 1:length(CC), b = 1:length(CC) ] - svdC = svd(G) - rk = rank(Diagonal(svdC.S), rtol = 1e-7) - # Diagonal(sqrt.(svdC.S[1:rk])) * svdC.U[:, 1:rk]' * CC - # construct the new basis - Ured = Diagonal(sqrt.(svdC.S[1:rk])) * svdC.U[:, 1:rk]' - Ure = Matrix{TCC}(undef, rk, length(Mll)) - for i = 1:rk - Ure[i, :] = sum(Ured[i, j] * CC[j] for j = 1:length(CC)) - end - return Ure, Mll + TP = typeof(A.phi) + TCC = coco_type(TP) + CC, Mll = compute_Al(A, ll) # CC::Vector{Vector{...}} + G = [sum(coco_dot(CC[a][i], CC[b][i]) for i = 1:length(Mll)) + for a = 1:length(CC), b = 1:length(CC)] + svdC = svd(G) + rk = rank(Diagonal(svdC.S), rtol=1e-7) + # Diagonal(sqrt.(svdC.S[1:rk])) * svdC.U[:, 1:rk]' * CC + # construct the new basis + Ured = Diagonal(sqrt.(svdC.S[1:rk])) * svdC.U[:, 1:rk]' + Ure = Matrix{TCC}(undef, rk, length(Mll)) + for i = 1:rk + Ure[i, :] = sum(Ured[i, j] * CC[j] for j = 1:length(CC)) + end + return Ure, Mll end # function barrier function compute_Al(A::Rot3DCoeffs, ll::SVector) - Mll = collect(_mrange(A.phi, ll)) - TP = typeof(A.phi) - if length(Mll) == 0 - return Vector{TP}[], Mll - end - - TA = typeof(A(ll, Mll[1], Mll[1])) - return __compute_Al(A, ll, Mll, TP, TA) + Mll = collect(_mrange(A.phi, ll)) + TP = typeof(A.phi) + if length(Mll) == 0 + return Vector{TP}[], Mll + end + + TA = typeof(A(ll, Mll[1], Mll[1])) + return __compute_Al(A, ll, Mll, TP, TA) end -# TODO: what was TA for? Can we get rid of it via coco_type? +# TODO: what was TA for? Can we get rid of it via coco_type? function __compute_Al(A::Rot3DCoeffs{T}, ll, Mll, TP, TA) where {T} - lenMll = length(Mll) - # each element of CC will be one row of the coupling coefficients - TCC = coco_type(TP) - CC = Vector{TCC}[] - # some utility funcions to allow coco_init to return either a property - # or a vector of properties - function __into_cc!(cc, cc0::AbstractProperty, im) - @assert length(cc) == 1 - cc[1][im] = cc0 - end - function __into_cc!(cc, cc0::AbstractVector, im) - @assert length(cc) == length(cc0) - for p = 1:length(cc) - cc[p][im] = cc0[p] - end - end - - for (ik, kk) in enumerate(Mll) # loop over possible basis functions - # do a dummy calculation to determine how many coefficients we will get - cc0 = A(ll, Mll[1], kk)::TA - numcc = (cc0 isa AbstractProperty ? 1 : length(cc0)) - # allocate the right number of vectors to store basis function coeffs - cc = [ Vector{TCC}(undef, lenMll) for _=1:numcc ] - for (im, mm) in enumerate(Mll) # loop over possible indices - if !coco_filter(A.phi, ll, mm, kk) - cc00 = zeros(TP, length(cc))::TA - __into_cc!(cc, cc00, im) - else - # get all possible coupling coefficients - cc0 = A(ll, mm, kk)::TA - __into_cc!(cc, cc0, im) - end - end - # and now push them onto the big stack. - append!(CC, cc) - end - - return CC, Mll + lenMll = length(Mll) + # each element of CC will be one row of the coupling coefficients + TCC = coco_type(TP) + CC = Vector{TCC}[] + # some utility funcions to allow coco_init to return either a property + # or a vector of properties + function __into_cc!(cc, cc0::AbstractProperty, im) + @assert length(cc) == 1 + cc[1][im] = cc0 + end + function __into_cc!(cc, cc0::AbstractVector, im) + @assert length(cc) == length(cc0) + for p = 1:length(cc) + cc[p][im] = cc0[p] + end + end + + for (ik, kk) in enumerate(Mll) # loop over possible basis functions + # do a dummy calculation to determine how many coefficients we will get + cc0 = A(ll, Mll[1], kk)::TA + numcc = (cc0 isa AbstractProperty ? 1 : length(cc0)) + # allocate the right number of vectors to store basis function coeffs + cc = [Vector{TCC}(undef, lenMll) for _ = 1:numcc] + for (im, mm) in enumerate(Mll) # loop over possible indices + if !coco_filter(A.phi, ll, mm, kk) + cc00 = zeros(TP, length(cc))::TA + __into_cc!(cc, cc00, im) + else + # get all possible coupling coefficients + cc0 = A(ll, mm, kk)::TA + __into_cc!(cc, cc0, im) + end + end + # and now push them onto the big stack. + append!(CC, cc) + end + + return CC, Mll end diff --git a/src/symmbasis.jl b/src/symmbasis.jl index d6200ef..86db977 100644 --- a/src/symmbasis.jl +++ b/src/symmbasis.jl @@ -26,145 +26,145 @@ Option 1: pass a `PIBasis` SymmetricBasis(φ, symgrp, pibasis) SymmetricBasis(φ, pibasis) ``` -If the PIbasis is already available, this directly constructs a -resulting SymmetricBasis; all possible permutation-invariant basis functions +If the PIbasis is already available, this directly constructs a +resulting SymmetricBasis; all possible permutation-invariant basis functions will be symmetrised and then reduced to a basis (rather than spanning set) """ -mutable struct SymmetricBasis{PIB, PROP, SYM, REAL} <: ACEBasis - pibasis::PIB - A2Bmap::SparseMatrixCSC{PROP, Int} - symgrp::SYM - real::REAL +mutable struct SymmetricBasis{PIB,PROP,SYM,REAL} <: ACEBasis + pibasis::PIB + A2Bmap::SparseMatrixCSC{PROP,Int} + symgrp::SYM + real::REAL end -Base.length(basis::SymmetricBasis{PIB, PROP}) where {PIB, PROP} = - size(basis.A2Bmap, 1) +Base.length(basis::SymmetricBasis{PIB,PROP}) where {PIB,PROP} = + size(basis.A2Bmap, 1) # -------- FIO -==(B1::SymmetricBasis, B2::SymmetricBasis) = - ( (B1.pibasis == B2.pibasis) && - (B1.A2Bmap == B2.A2Bmap) && - (B1.real == B2.real) ) - -write_dict(B::SymmetricBasis{PIB, PROP}) where {PIB, PROP} = - Dict( "__id__" => "ACEfrictionCore_SymmetricBasis", - "pibasis" => write_dict(B.pibasis), - "A2Bmap" => write_dict(B.A2Bmap), - "symgrp" => write_dict(B.symgrp), - "isreal" => (B.real == Base.real) ) - -read_dict(::Val{:ACEfrictionCore_SymmetricBasis}, D::Dict) = - SymmetricBasis(read_dict(D["pibasis"]), - read_dict(D["A2Bmap"]), - read_dict(D["symgrp"]), - (D["isreal"] ? Base.real : Base.identity) ) +==(B1::SymmetricBasis, B2::SymmetricBasis) = + ((B1.pibasis == B2.pibasis) && + (B1.A2Bmap == B2.A2Bmap) && + (B1.real == B2.real)) + +write_dict(B::SymmetricBasis{PIB,PROP}) where {PIB,PROP} = + Dict("__id__" => "ACEfrictionCore_SymmetricBasis", + "pibasis" => write_dict(B.pibasis), + "A2Bmap" => write_dict(B.A2Bmap), + "symgrp" => write_dict(B.symgrp), + "isreal" => (B.real == Base.real)) + +read_dict(::Val{:ACEfrictionCore_SymmetricBasis}, D::AbstractDict) = + SymmetricBasis(read_dict(D["pibasis"]), + read_dict(D["A2Bmap"]), + read_dict(D["symgrp"]), + (D["isreal"] ? Base.real : Base.identity)) # -------- -SymmetricBasis(φ::AbstractProperty, - basis1p::OneParticleBasis, - Bsel::AbstractBasisSelector; - kwargs...) = - SymmetricBasis(φ, basis1p, O3(), Bsel; kwargs...) - -SymmetricBasis(φ::AbstractProperty, pibasis; kwargs...) = - SymmetricBasis(φ, O3(), pibasis; kwargs...) - - -SymmetricBasis(φ::AbstractProperty, - basis1p::OneParticleBasis, - symgrp::SymmetryGroup, - Bsel::AbstractBasisSelector; - isreal=isrealB(φ), kwargs...) = - SymmetricBasis(φ, symgrp, - PIBasis(basis1p, symgrp, Bsel; - isreal = isrealAA(φ), kwargs..., property = φ); - isreal=isreal) - -SymmetricBasis(φ::AbstractProperty, symgrp::SymmetryGroup, pibasis::PIBasis; - isreal=false) = - SymmetricBasis(φ, symgrp, pibasis, isreal ? Base.real : Base.identity) - - -function SymmetricBasis(φ::TP, symgrp::SymmetryGroup, pibasis::PIBasis, - _real) where {TP <: AbstractProperty} - - # AA index -> AA spec - AAspec = get_spec(pibasis) - - # construct the reverse mapping AA spec -> iAA and A spec -> iA - invAAspec = Dict{Any, Int}() - for (iAA, AA) in enumerate(AAspec) - invAAspec[AA] = iAA - end - invAspec = Dict{Any, Int}() - for (iA, A) in enumerate(get_spec(pibasis.basis1p)) - invAspec[A] = iA - end - - # allocate the datastructure that computes and caches the - # coupling coefficients - # TODO: should this be stored with the basis? - # or maybe written to a file on disk? and then flushed every time - # we finish with a basis construction??? - # TODO: for sure this needs to become a function of the symmetry group? - rotc = Rot3DCoeffs(φ, Float64) - # allocate triplet format - TCC = coco_type(TP) - Irow, Jcol, vals = Int[], Int[], TCC[] - # count the number of PI basis functions = number of rows - idxB = 0 - - # loop through AA basis, but skip most of them ... - for (iAA, AA) in enumerate(AAspec) - # determine whether we need to compute coupling coefficients for this - # basis function or whether it will be included in a different - # coco computation? - if !is_refbasisfcn(symgrp, AA) - continue - end - # compute the cocos - U, AAcols = coupling_coeffs(symgrp, AA, rotc) - - # loop over the rows of U -> each specifies a basis function which we now - # need to incorporate into the basis - for irow = 1:size(U, 1) - idxB += 1 - # loop over the columns of U / over brows - for (icol, bcol) in enumerate(AAcols) - # put bcol into the correct order - bcol_ordered = _get_ordered(bcol, invAspec) - # this is a subtle step: bcol and bcol_ordered are equivalent - # permutation-invariant basis functions. This means we will - # add the same PI basis function several times, but in the call to - # `sparse` the values will just be added. - if !haskey(invAAspec, bcol_ordered) - @warn("bcol_ordered not in AA-spec") - @show bcol_ordered - # @show degree(bcol_ordered, NaiveTotalDegree(), pibasis.basis1p) - error("bcol_ordered not in AA-spec") +SymmetricBasis(φ::AbstractProperty, + basis1p::OneParticleBasis, + Bsel::AbstractBasisSelector; + kwargs...) = + SymmetricBasis(φ, basis1p, O3(), Bsel; kwargs...) + +SymmetricBasis(φ::AbstractProperty, pibasis; kwargs...) = + SymmetricBasis(φ, O3(), pibasis; kwargs...) + + +SymmetricBasis(φ::AbstractProperty, + basis1p::OneParticleBasis, + symgrp::SymmetryGroup, + Bsel::AbstractBasisSelector; + isreal=isrealB(φ), kwargs...) = + SymmetricBasis(φ, symgrp, + PIBasis(basis1p, symgrp, Bsel; + isreal=isrealAA(φ), kwargs..., property=φ); + isreal=isreal) + +SymmetricBasis(φ::AbstractProperty, symgrp::SymmetryGroup, pibasis::PIBasis; + isreal=false) = + SymmetricBasis(φ, symgrp, pibasis, isreal ? Base.real : Base.identity) + + +function SymmetricBasis(φ::TP, symgrp::SymmetryGroup, pibasis::PIBasis, + _real) where {TP<:AbstractProperty} + + # AA index -> AA spec + AAspec = get_spec(pibasis) + + # construct the reverse mapping AA spec -> iAA and A spec -> iA + invAAspec = Dict{Any,Int}() + for (iAA, AA) in enumerate(AAspec) + invAAspec[AA] = iAA + end + invAspec = Dict{Any,Int}() + for (iA, A) in enumerate(get_spec(pibasis.basis1p)) + invAspec[A] = iA + end + + # allocate the datastructure that computes and caches the + # coupling coefficients + # TODO: should this be stored with the basis? + # or maybe written to a file on disk? and then flushed every time + # we finish with a basis construction??? + # TODO: for sure this needs to become a function of the symmetry group? + rotc = Rot3DCoeffs(φ, Float64) + # allocate triplet format + TCC = coco_type(TP) + Irow, Jcol, vals = Int[], Int[], TCC[] + # count the number of PI basis functions = number of rows + idxB = 0 + + # loop through AA basis, but skip most of them ... + for (iAA, AA) in enumerate(AAspec) + # determine whether we need to compute coupling coefficients for this + # basis function or whether it will be included in a different + # coco computation? + if !is_refbasisfcn(symgrp, AA) + continue + end + # compute the cocos + U, AAcols = coupling_coeffs(symgrp, AA, rotc) + + # loop over the rows of U -> each specifies a basis function which we now + # need to incorporate into the basis + for irow = 1:size(U, 1) + idxB += 1 + # loop over the columns of U / over brows + for (icol, bcol) in enumerate(AAcols) + # put bcol into the correct order + bcol_ordered = _get_ordered(bcol, invAspec) + # this is a subtle step: bcol and bcol_ordered are equivalent + # permutation-invariant basis functions. This means we will + # add the same PI basis function several times, but in the call to + # `sparse` the values will just be added. + if !haskey(invAAspec, bcol_ordered) + @warn("bcol_ordered not in AA-spec") + @show bcol_ordered + # @show degree(bcol_ordered, NaiveTotalDegree(), pibasis.basis1p) + error("bcol_ordered not in AA-spec") + end + idxAA = invAAspec[bcol_ordered] + push!(Irow, idxB) + push!(Jcol, idxAA) + push!(vals, U[irow, icol]) end - idxAA = invAAspec[bcol_ordered] - push!(Irow, idxB) - push!(Jcol, idxAA) - push!(vals, U[irow, icol]) - end - end - end - # TODO: filter and throw out everything that hasn't been used!! - # create CSC: [ triplet ] nrows ncols - A2Bmap = sparse(Irow, Jcol, vals, idxB, length(AAspec)) - basis = SymmetricBasis(pibasis, A2Bmap, symgrp, _real) - # clean up a bit, i.e. remove AA basis functions that we don't need - # to evaluate the symmetric basis - clean_pibasis!(basis) - return basis + end + end + # TODO: filter and throw out everything that hasn't been used!! + # create CSC: [ triplet ] nrows ncols + A2Bmap = sparse(Irow, Jcol, vals, idxB, length(AAspec)) + basis = SymmetricBasis(pibasis, A2Bmap, symgrp, _real) + # clean up a bit, i.e. remove AA basis functions that we don't need + # to evaluate the symmetric basis + clean_pibasis!(basis) + return basis end -function SymmetricBasis(pibasis, A2Bmap, symgrp, _real) - PROP = _real(eltype(A2Bmap)) - B_pool = VectorPool{PROP}() - return SymmetricBasis(pibasis, A2Bmap, symgrp, _real, B_pool) +function SymmetricBasis(pibasis, A2Bmap, symgrp, _real) + PROP = _real(eltype(A2Bmap)) + B_pool = VectorPool{PROP}() + return SymmetricBasis(pibasis, A2Bmap, symgrp, _real, B_pool) end @@ -172,75 +172,75 @@ end produce the ordered tuple defining the AA basis function uniquely. """ function _get_ordered(bb, invAspec) - iAs = [invAspec[b] for b in bb] - return bb[ sortperm(iAs, rev = true) ] + iAs = [invAspec[b] for b in bb] + return bb[sortperm(iAs, rev=true)] end -# ------------------- exporting the basis spec +# ------------------- exporting the basis spec -# this doesn't provide the "full" specification, just collects -# the n and l but not the m or coupling coefficients. +# this doesn't provide the "full" specification, just collects +# the n and l but not the m or coupling coefficients. function get_spec(basis::SymmetricBasis) - spec_AA = get_spec(basis.pibasis) - spec_B = [] - for iB = 1:length(basis) - iAA = findfirst(norm.(basis.A2Bmap[iB, :]) .!= 0) - push!(spec_B, get_sym_spec(basis.symgrp, spec_AA[iAA])) - end - return identity.(spec_B) + spec_AA = get_spec(basis.pibasis) + spec_B = [] + for iB = 1:length(basis) + iAA = findfirst(norm.(basis.A2Bmap[iB, :]) .!= 0) + push!(spec_B, get_sym_spec(basis.symgrp, spec_AA[iAA])) + end + return identity.(spec_B) end -# ---------------- sparsification and cleanup codes +# ---------------- sparsification and cleanup codes """ `sparsify!(basis::SymmetricBasis; del = ..., keep = ...)` -sparsify the symmetric basis by either specifyin which basis functions to +sparsify the symmetric basis by either specifyin which basis functions to keep or which ones to delete. """ -function sparsify!(basis::SymmetricBasis; - del::Union{Nothing, AbstractVector{<: Integer}} = nothing, - keep::Union{Nothing, AbstractVector{<: Integer}} = nothing) - if ( (del == nothing && keep == nothing ) || - (del != nothing && keep != nothing ) ) - error("sparsify!: must provide either del or keep kwarg but not both") - end - if del != nothing - keep = setdiff(1:length(basis), del) - end - basis.A2Bmap = basis.A2Bmap[keep, :] - clean_pibasis!(basis; atol = 0.0) +function sparsify!(basis::SymmetricBasis; + del::Union{Nothing,AbstractVector{<:Integer}}=nothing, + keep::Union{Nothing,AbstractVector{<:Integer}}=nothing) + if ((del == nothing && keep == nothing) || + (del != nothing && keep != nothing)) + error("sparsify!: must provide either del or keep kwarg but not both") + end + if del != nothing + keep = setdiff(1:length(basis), del) + end + basis.A2Bmap = basis.A2Bmap[keep, :] + clean_pibasis!(basis; atol=0.0) end # this is also used for Filtering the AA basis if there are zero-rows in the A2B map """ Remove elements of the AA basis, when there are zero-rows in the A2B map, i.e. -when some AA basis elements are simply not required to evaluate the -symmetric basis. +when some AA basis elements are simply not required to evaluate the +symmetric basis. """ -function clean_pibasis!(basis::SymmetricBasis; atol = 0.0) - # get the zero-columns - Inz = sort( findall(x -> x > atol, sum(norm, basis.A2Bmap, dims=1)[:]) ) - if length(Inz) < size(basis.A2Bmap, 2) - # remove those columns from the A2Bmap - sparsify!(basis.pibasis, Inz) - # remove those columns from the A2Bmap - basis.A2Bmap = basis.A2Bmap[:, Inz] - end - clean_1pbasis!(basis.pibasis) - return basis +function clean_pibasis!(basis::SymmetricBasis; atol=0.0) + # get the zero-columns + Inz = sort(findall(x -> x > atol, sum(norm, basis.A2Bmap, dims=1)[:])) + if length(Inz) < size(basis.A2Bmap, 2) + # remove those columns from the A2Bmap + sparsify!(basis.pibasis, Inz) + # remove those columns from the A2Bmap + basis.A2Bmap = basis.A2Bmap[:, Inz] + end + clean_1pbasis!(basis.pibasis) + return basis end # ---------------- A modified sparse matmul -# TODO: move this stuff all to aux? +# TODO: move this stuff all to aux? using SparseArrays: AbstractSparseMatrixCSC, - nonzeros, rowvals, nzrange + nonzeros, rowvals, nzrange using LinearAlgebra: Transpose @@ -254,7 +254,7 @@ function genmul!(C, A::AbstractSparseMatrixCSC, B, mulop) fill!(C, zero(eltype(C))) for k in 1:size(C, 2) @inbounds for col in 1:size(A, 2) - αxj = B[col,k] + αxj = B[col, k] for j in nzrange(A, col) C[rv[j], k] += mulop(nzv[j], αxj) end @@ -265,66 +265,64 @@ end function genmul!(C, xA::Transpose{<:Any,<:AbstractSparseMatrixCSC}, B, mulop) - A = xA.parent - size(A, 2) == size(C, 1) || throw(DimensionMismatch()) - size(A, 1) == size(B, 1) || throw(DimensionMismatch()) - size(B, 2) == size(C, 2) || throw(DimensionMismatch()) - nzv = nonzeros(A) - rv = rowvals(A) - fill!(C, zero(eltype(C))) - for k in 1:size(C, 2) - @inbounds for col in 1:size(A, 2) - tmp = zero(eltype(C)) - for j in nzrange(A, col) - tmp += mulop(nzv[j], B[rv[j],k]) - end - C[col,k] += tmp - end - end - return C + A = xA.parent + size(A, 2) == size(C, 1) || throw(DimensionMismatch()) + size(A, 1) == size(B, 1) || throw(DimensionMismatch()) + size(B, 2) == size(C, 2) || throw(DimensionMismatch()) + nzv = nonzeros(A) + rv = rowvals(A) + fill!(C, zero(eltype(C))) + for k in 1:size(C, 2) + @inbounds for col in 1:size(A, 2) + tmp = zero(eltype(C)) + for j in nzrange(A, col) + tmp += mulop(nzv[j], B[rv[j], k]) + end + C[col, k] += tmp + end + end + return C end function genmul(A, B, mulop) - T = typeof(mulop(A[1], B[1])) - C = Array{T}(undef, (size(A, 1), size(B)[2:end]...) ) - return genmul!(C, A, B, mulop) + T = typeof(mulop(A[1], B[1])) + C = Array{T}(undef, (size(A, 1), size(B)[2:end]...)) + return genmul!(C, A, B, mulop) end # ---------------- Evaluation code function evaluate(basis::SymmetricBasis, cfg::UConfig) - AA = evaluate(basis.pibasis, cfg) - B = evaluate(basis, AA) - release!(AA) - return B + AA = evaluate(basis.pibasis, cfg) + B = evaluate(basis, AA) + release!(AA) + return B end -# NOTE: Nasty and completely not understood type instability here +# NOTE: Nasty and completely not understood type instability here function evaluate!(B, basis::SymmetricBasis, cfg::UConfig) - AA = evaluate(basis.pibasis, cfg) - evaluate!(B, basis, AA) - release!(AA) - return B + AA = evaluate(basis.pibasis, cfg) + evaluate!(B, basis, AA) + release!(AA) + return B end -evaluate(basis::SymmetricBasis, AA::AbstractVector{<: Number}) = - genmul(basis.A2Bmap, AA, (a, b) -> basis.real(a * b)) +evaluate(basis::SymmetricBasis, AA::AbstractVector{<:Number}) = + genmul(basis.A2Bmap, AA, (a, b) -> basis.real(a * b)) -evaluate!(B, basis::SymmetricBasis, AA::AbstractVector{<: Number}) = - genmul!(B, basis.A2Bmap, AA, (a, b) -> basis.real(a * b)) +evaluate!(B, basis::SymmetricBasis, AA::AbstractVector{<:Number}) = + genmul!(B, basis.A2Bmap, AA, (a, b) -> basis.real(a * b)) # ---------------- gradients -# ------------------------------- +# ------------------------------- function scaling(basis::SymmetricBasis, p) - wwpi = scaling(basis.pibasis, p) - wwrpi = abs2.(norm.(basis.A2Bmap)) * abs2.(wwpi) - return sqrt.(wwrpi) + wwpi = scaling(basis.pibasis, p) + wwrpi = abs2.(norm.(basis.A2Bmap)) * abs2.(wwpi) + return sqrt.(wwrpi) end - - diff --git a/src/symmetrygroups.jl b/src/symmetrygroups.jl index ad31405..5cb9bb0 100644 --- a/src/symmetrygroups.jl +++ b/src/symmetrygroups.jl @@ -1,324 +1,325 @@ -# NOTE: at the moment these are ad hoc implementations for each group that -# we decide we need. eventually we can hopefully simplify and merge +# NOTE: at the moment these are ad hoc implementations for each group that +# we decide we need. eventually we can hopefully simplify and merge # many of these codes -# The first version was written before NamedTupleTools (or before I +# The first version was written before NamedTupleTools (or before I # knew about it -> this will simplify the code here a bit) using NamedTupleTools: delete, merge, namedtuple -abstract type SymmetryGroup end +abstract type SymmetryGroup end """ -`struct NoSym <: SymmetryGroup ` : no symmetrisation other than +`struct NoSym <: SymmetryGroup ` : no symmetrisation other than permutation symmetry already baked into the AA basis. """ """ -`struct NoSym <: SymmetryGroup` : no symmetry beyond the standard -permutation symmetry. This is currently not used, but could be incorporated -to provide a more streamlined experience for the user. +`struct NoSym <: SymmetryGroup` : no symmetry beyond the standard +permutation symmetry. This is currently not used, but could be incorporated +to provide a more streamlined experience for the user. """ -struct NoSym <: SymmetryGroup -end +struct NoSym <: SymmetryGroup +end -# write_dict(G::O3) = -# Dict("__id__" => "ACEfrictionCore_O3", -# "lsym" => lsym(G), +# write_dict(G::O3) = +# Dict("__id__" => "ACEfrictionCore_O3", +# "lsym" => lsym(G), # "msym" => msym(G) ) -# read_dict(::Val{:ACEfrictionCore_O3}, D::Dict) = +# read_dict(::Val{:ACEfrictionCore_O3}, D::AbstractDict) = # O3(Symbol(D["lsym"]), Symbol(D["msym"])) -is_refbasisfcn(G::NoSym, AA) = true +is_refbasisfcn(G::NoSym, AA) = true get_sym_spec(G::NoSym, bb) = bb -function coupling_coeffs(symgrp::NoSym, bb, rotc::Rot3DCoeffs{T, TP}) where {T, TP} - return [ one(TP) ], [bb] +function coupling_coeffs(symgrp::NoSym, bb, rotc::Rot3DCoeffs{T,TP}) where {T,TP} + return [one(TP)], [bb] end -# ----------------------- O3 SYMMETRY +# ----------------------- O3 SYMMETRY -# this is a prototype implemenation; eventually (asap!) we need to allow +# this is a prototype implemenation; eventually (asap!) we need to allow # rotation of multiple features at once, e.g., spin-orbit coupling! """ -`struct O3 <: SymmetryGroup` : this is the default symmetry group; describing -the action of a single O3 group on the basis. +`struct O3 <: SymmetryGroup` : this is the default symmetry group; describing +the action of a single O3 group on the basis. -Standard Usage: -```julia +Standard Usage: +```julia O3() ``` -will create an `O3{:l, :m}` instance, i.e. the group will expect the symbols -`:l, :m` in the relevant 1p basis. - -But if the Ylm component of the 1p basis uses different symbols then one can -tell `O3` this via `O3(lsym, rsym)`. E.g. if the variable w.r.t. which we -symmetrize is a spin `s` then we might call it `O3(:ls, :ms)`. The main thing -to remember is that the symbols in the Ylm basis and in the O3 basis must -match. +will create an `O3{:l, :m}` instance, i.e. the group will expect the symbols +`:l, :m` in the relevant 1p basis. + +But if the Ylm component of the 1p basis uses different symbols then one can +tell `O3` this via `O3(lsym, rsym)`. E.g. if the variable w.r.t. which we +symmetrize is a spin `s` then we might call it `O3(:ls, :ms)`. The main thing +to remember is that the symbols in the Ylm basis and in the O3 basis must +match. """ -struct O3{LSYM, MSYM} <: SymmetryGroup +struct O3{LSYM,MSYM} <: SymmetryGroup end -O3(lsym::Symbol = :l, msym::Symbol = :m) = O3{lsym, msym}() +O3(lsym::Symbol=:l, msym::Symbol=:m) = O3{lsym,msym}() -lsym(G::O3{LSYM, MSYM}) where {LSYM, MSYM} = LSYM +lsym(G::O3{LSYM,MSYM}) where {LSYM,MSYM} = LSYM getl(G::O3, b::NamedTuple) = b[lsym(G)] -msym(G::O3{LSYM, MSYM}) where {LSYM, MSYM} = MSYM +msym(G::O3{LSYM,MSYM}) where {LSYM,MSYM} = MSYM getm(G::O3, b::NamedTuple) = b[msym(G)] -write_dict(G::O3) = - Dict("__id__" => "ACEfrictionCore_O3", - "lsym" => lsym(G), - "msym" => msym(G) ) +write_dict(G::O3) = + Dict("__id__" => "ACEfrictionCore_O3", + "lsym" => lsym(G), + "msym" => msym(G)) -read_dict(::Val{:ACEfrictionCore_O3}, D::Dict) = - O3(Symbol(D["lsym"]), Symbol(D["msym"])) +read_dict(::Val{:ACEfrictionCore_O3}, D::AbstractDict) = + O3(Symbol(D["lsym"]), Symbol(D["msym"])) -is_refbasisfcn(G::O3, AA) = all( bi[msym(G)] == 0 for bi in AA ) +is_refbasisfcn(G::O3, AA) = all(bi[msym(G)] == 0 for bi in AA) get_sym_spec(G::O3, bb) = delete.(bb, (msym(G),)) -function coupling_coeffs(symgrp::O3, bb, rotc::Rot3DCoeffs{T, TP}) where {T, TP} - # bb = [ b1, b2, b3, ... ] - # bi = (μ = ..., n = ..., l = ..., m = ...) - # (μ, n) -> n; only the l and m are used in the angular basis - if length(bb) == 0 - return coco_init(rotc.phi), [bb,] - end - # convert to a format that the Rotations3D implementation can understand - # this utility function splits the bb = (b1, b2 ...) with each - # b1 = (μ = ..., n = ..., l = ..., m = ...) into - # l, and a new n = (μ, n) - ll, nn = _b2llnn(symgrp, bb) - # now we can call the coupling coefficient construiction!! - U, Ms = rpe_basis(rotc, nn, ll) - - # but now we need to convert the m spec back to complete basis function - # specifications (provided by sending in a prototype b = bb[1]) - rpebs = [ _nnllmm2b(symgrp, bb[1], nn, ll, mm) for mm in Ms ] - - return U, rpebs +function coupling_coeffs(symgrp::O3, bb, rotc::Rot3DCoeffs{T,TP}) where {T,TP} + # bb = [ b1, b2, b3, ... ] + # bi = (μ = ..., n = ..., l = ..., m = ...) + # (μ, n) -> n; only the l and m are used in the angular basis + if length(bb) == 0 + return coco_init(rotc.phi), [bb,] + end + # convert to a format that the Rotations3D implementation can understand + # this utility function splits the bb = (b1, b2 ...) with each + # b1 = (μ = ..., n = ..., l = ..., m = ...) into + # l, and a new n = (μ, n) + ll, nn = _b2llnn(symgrp, bb) + # now we can call the coupling coefficient construiction!! + U, Ms = rpe_basis(rotc, nn, ll) + + # but now we need to convert the m spec back to complete basis function + # specifications (provided by sending in a prototype b = bb[1]) + rpebs = [_nnllmm2b(symgrp, bb[1], nn, ll, mm) for mm in Ms] + + return U, rpebs end function rpe_basis(A::Rot3DCoeffs, - nn::SVector{N, TN}, - ll::SVector{N, Int}) where {N, TN} - Ure, Mre = Rotations3D.re_basis(A, ll) - G = _gramian(nn, ll, Ure, Mre) - S = svd(G) - rk = rank(Diagonal(S.S); rtol = 1e-7) - Urpe = S.U[:, 1:rk]' - return Diagonal(sqrt.(S.S[1:rk])) * Urpe * Ure, Mre + nn::SVector{N,TN}, + ll::SVector{N,Int}) where {N,TN} + Ure, Mre = Rotations3D.re_basis(A, ll) + G = _gramian(nn, ll, Ure, Mre) + S = svd(G) + rk = rank(Diagonal(S.S); rtol=1e-7) + Urpe = S.U[:, 1:rk]' + return Diagonal(sqrt.(S.S[1:rk])) * Urpe * Ure, Mre end function _gramian(nn, ll, Ure, Mre) - N = length(nn) - nre = size(Ure, 1) - G = zeros(Complex{Float64}, nre, nre) - for σ in permutations(1:N) - if (nn[σ] != nn) || (ll[σ] != ll); continue; end - for (iU1, mm1) in enumerate(Mre), (iU2, mm2) in enumerate(Mre) - if mm1[σ] == mm2 - for i1 = 1:nre, i2 = 1:nre - G[i1, i2] += coco_dot(Ure[i1, iU1], Ure[i2, iU2]) + N = length(nn) + nre = size(Ure, 1) + G = zeros(Complex{Float64}, nre, nre) + for σ in permutations(1:N) + if (nn[σ] != nn) || (ll[σ] != ll) + continue + end + for (iU1, mm1) in enumerate(Mre), (iU2, mm2) in enumerate(Mre) + if mm1[σ] == mm2 + for i1 = 1:nre, i2 = 1:nre + G[i1, i2] += coco_dot(Ure[i1, iU1], Ure[i2, iU2]) + end end - end - end - end - return G + end + end + return G end # TODO: replace all this awful code with NamedTupleTools -_nnllmm2b(G, b, nn, ll, mm) = [ _nlm2b(G, b, n, l, m) for (n, l, m) in zip(nn, ll, mm) ] +_nnllmm2b(G, b, nn, ll, mm) = [_nlm2b(G, b, n, l, m) for (n, l, m) in zip(nn, ll, mm)] -@generated function _nlm2b(G::O3{LSYM, MSYM}, b::NamedTuple{ALLKEYS}, - n::NamedTuple{NKEYS}, - l, m) where {LSYM, MSYM, ALLKEYS, NKEYS} - code = - ( "( _b = (" * prod("$(k) = n.$(k), " for k in NKEYS) - * "$(LSYM) = l, $(MSYM) = m ); " +@generated function _nlm2b(G::O3{LSYM,MSYM}, b::NamedTuple{ALLKEYS}, + n::NamedTuple{NKEYS}, + l, m) where {LSYM,MSYM,ALLKEYS,NKEYS} + code = + ("( _b = (" * prod("$(k) = n.$(k), " for k in NKEYS) + * "$(LSYM) = l, $(MSYM) = m ); " * - " b = (" * prod("$(k) = _b.$(k), " for k in ALLKEYS) * ") )" ) - :( $(Meta.parse(code)) ) + " b = (" * prod("$(k) = _b.$(k), " for k in ALLKEYS) * ") )") + :($(Meta.parse(code))) end function _b2llnn(G::O3, bb) - @assert all( iszero(b[msym(G)]) for b in bb ) - ll = SVector( [b[lsym(G)] for b in bb]... ) - nn = SVector( [_all_but_lm(G, b) for b in bb]... ) - return ll, nn + @assert all(iszero(b[msym(G)]) for b in bb) + ll = SVector([b[lsym(G)] for b in bb]...) + nn = SVector([_all_but_lm(G, b) for b in bb]...) + return ll, nn end """ return a NamedTuple containing all values in b except those corresponding to l and m keys -TODO: get rid of this ridiculousness and replace with NamedTupleTools methods +TODO: get rid of this ridiculousness and replace with NamedTupleTools methods """ -@generated function _all_but_lm(G::O3{LSYM, MSYM}, - b::NamedTuple{NAMES}) where {LSYM, MSYM, NAMES} - code = "n = (" - for k in NAMES - if !(k in (LSYM, MSYM)) - code *= "$(k) = b.$(k), " - end - end - code *= ")" - quote - $(Meta.parse(code)) - n - end +@generated function _all_but_lm(G::O3{LSYM,MSYM}, + b::NamedTuple{NAMES}) where {LSYM,MSYM,NAMES} + code = "n = (" + for k in NAMES + if !(k in (LSYM, MSYM)) + code *= "$(k) = b.$(k), " + end + end + code *= ")" + quote + $(Meta.parse(code)) + n + end end -# -------------- O3 ⊗ O3 +# -------------- O3 ⊗ O3 -# this is a preliminary implementation; eventually we may want a more -# general description composition of arbitrary isometry combinations +# this is a preliminary implementation; eventually we may want a more +# general description composition of arbitrary isometry combinations @doc raw""" -`struct O3O3 <: SymmetryGroup` : This type implements the ``O(3) \otimes O(3)`` symmetry -group. This is useful when a particle has two euclidean vector attributes, say -``{\bm r}`` and ``{\bm s}`` and the action of the group on the pair is -```math +`struct O3O3 <: SymmetryGroup` : This type implements the ``O(3) \otimes O(3)`` symmetry +group. This is useful when a particle has two euclidean vector attributes, say +``{\bm r}`` and ``{\bm s}`` and the action of the group on the pair is +```math (Q_r, Q_s)[ (\boldsymbol{r}, \boldsymbol{s}) ] = (Q_r \boldsymbol{r}, Q_s \boldsymbol{s}) ``` -A canA canonical application is magnetism: it is known that spin-orbit coupling -is a very weak effect. By ignoring it, i.e., letting positions and spins rotate -independently of one another, one makes a small modelling error. This leads -precisely to the ``O(3) \otimes O(3)`` symmetry. +A canA canonical application is magnetism: it is known that spin-orbit coupling +is a very weak effect. By ignoring it, i.e., letting positions and spins rotate +independently of one another, one makes a small modelling error. This leads +precisely to the ``O(3) \otimes O(3)`` symmetry. -To construct this group, use +To construct this group, use ```julia symgrp = O3(:lr, :mr) ⊗ O3(:ls, :ms) ``` -or replace those symbols with the appropriate symbols used to specify the -corresponding `Ylm1pbasis` objects. +or replace those symbols with the appropriate symbols used to specify the +corresponding `Ylm1pbasis` objects. """ -struct O3O3{LSYM1, MSYM1, LSYM2, MSYM2} <: SymmetryGroup - G1::O3{LSYM1, MSYM1} - G2::O3{LSYM2, MSYM2} +struct O3O3{LSYM1,MSYM1,LSYM2,MSYM2} <: SymmetryGroup + G1::O3{LSYM1,MSYM1} + G2::O3{LSYM2,MSYM2} end import Base: kron function kron(G1::O3, G2::O3) - @assert lsym(G1) != lsym(G2) - @assert msym(G1) != msym(G2) - return O3O3(G1, G2) + @assert lsym(G1) != lsym(G2) + @assert msym(G1) != msym(G2) + return O3O3(G1, G2) end ⊗(G1::O3, G2::O3) = kron(G1, G2) export ⊗ -write_dict(G::O3O3) = - Dict("__id__" => "ACEfrictionCore_O3O3", - "G1" => write_dict(G.G1), - "G2" => write_dict(G.G2) ) +write_dict(G::O3O3) = + Dict("__id__" => "ACEfrictionCore_O3O3", + "G1" => write_dict(G.G1), + "G2" => write_dict(G.G2)) -read_dict(::Val{:ACEfrictionCore_O3O3}, D::Dict) = - read_dict(D["G1"]) ⊗ read_dict(D["G2"]) +read_dict(::Val{:ACEfrictionCore_O3O3}, D::AbstractDict) = + read_dict(D["G1"]) ⊗ read_dict(D["G2"]) -is_refbasisfcn(G::O3O3, AA) = all( bi[msym(grp)] == 0 - for bi in AA, grp in (G.G1, G.G2) ) +is_refbasisfcn(G::O3O3, AA) = all(bi[msym(grp)] == 0 + for bi in AA, grp in (G.G1, G.G2)) -get_sym_spec(G::O3O3, bb) = delete.(bb, Ref( (msym(G.G1), msym(G.G2)) )) +get_sym_spec(G::O3O3, bb) = delete.(bb, Ref((msym(G.G1), msym(G.G2)))) function coupling_coeffs(symgrp::O3O3, bb, rotc::Rot3DCoeffs) - # bb = [ b1, b2, b3, ... ] - # bi = (μ = ..., n = ..., l1 = ..., m1 = ..., l2 = ..., m2 = ...) - # (μ, n, ...) -> n; only the l and m are used in the angular basis - if length(bb) == 0 - return coco_init(rotc.phi), [bb,] - end - - # the prototype namedtuple describing a single 1p basis fcn - PROTOTUPLE = prototype(bb[1]) - NU = length(bb) - - # convert to (nn, ll, mm) format for Rotations3D - ll1, ll2, nn, ll12 = _b2llnn(symgrp, bb) - # ... and construct the coupling coefficients for the individual subgroups - U1, M1 = Rotations3D.re_basis(rotc, ll1) - U2, M2 = Rotations3D.re_basis(rotc, ll2) - - nU1, nM1 = size(U1) - nU2, nM2 = size(U2) - @assert nM1 == length(M1) - @assert nM2 == length(M2) - UT = promote_type(eltype(U1), eltype(U2)) - - # there is admissible combination: - if nU1 == 0 || nU2 == 0 - return UT[], SVector{NU, PROTOTUPLE}[] - end - - # now combine them into the effective coupling coeffs and combined Ms - # each column Ure[:, i] corresponds to one rotation-invariant basis fcn - M1M2TUPLE = namedtuple(msym(symgrp.G1), msym(symgrp.G2)) - Mre = Vector{typeof(Vector(M1M2TUPLE.(M1[1], M2[1])))}(undef, nM1 * nM2) - Ure = zeros( UT, (nU1 * nU2, nM1 * nM2) ) - jdx = 0 - for j1 = 1:nM1, j2 = 1:nM2 - jdx += 1 - Mre[jdx] = Vector(M1M2TUPLE.(M1[j1], M2[j2])) - idx = 0 - for i1 = 1:nU1, i2 = 1:nU2 - idx += 1 - Ure[idx, jdx] = U1[i1, j1] * U2[i2, j2] - end - end - - # insert another reduction step - # in my tests this never reduces the size, but I haven't had the time - # to actually prove it isn't needed, so we will keep it for now - Gre = [ sum(coco_dot.(Ure[i1, :], Ure[i2, :])) for i1 = 1:size(Ure, 1), i2 = 1:size(Ure, 1) ] - Sre = svd(Gre) - rk = rank(Diagonal(Sre.S); rtol = 1e-7) - Ure = Sre.U[:, 1:rk]' * Ure - - # now symmetrize w.r.t. permutations - G = _gramian(nn, ll12, Ure, Mre) - S = svd(G) - rk = rank(Diagonal(S.S); rtol = 1e-7) - Urpe = S.U[:, 1:rk]' - U = Diagonal(sqrt.(S.S[1:rk])) * Urpe * Ure - - # reconstruct the basis function specifications - rpebs = [ PROTOTUPLE.(merge.(nn, ll12, mm12)) for mm12 in Mre ] - - return U, rpebs + # bb = [ b1, b2, b3, ... ] + # bi = (μ = ..., n = ..., l1 = ..., m1 = ..., l2 = ..., m2 = ...) + # (μ, n, ...) -> n; only the l and m are used in the angular basis + if length(bb) == 0 + return coco_init(rotc.phi), [bb,] + end + + # the prototype namedtuple describing a single 1p basis fcn + PROTOTUPLE = prototype(bb[1]) + NU = length(bb) + + # convert to (nn, ll, mm) format for Rotations3D + ll1, ll2, nn, ll12 = _b2llnn(symgrp, bb) + # ... and construct the coupling coefficients for the individual subgroups + U1, M1 = Rotations3D.re_basis(rotc, ll1) + U2, M2 = Rotations3D.re_basis(rotc, ll2) + + nU1, nM1 = size(U1) + nU2, nM2 = size(U2) + @assert nM1 == length(M1) + @assert nM2 == length(M2) + UT = promote_type(eltype(U1), eltype(U2)) + + # there is admissible combination: + if nU1 == 0 || nU2 == 0 + return UT[], SVector{NU,PROTOTUPLE}[] + end + + # now combine them into the effective coupling coeffs and combined Ms + # each column Ure[:, i] corresponds to one rotation-invariant basis fcn + M1M2TUPLE = namedtuple(msym(symgrp.G1), msym(symgrp.G2)) + Mre = Vector{typeof(Vector(M1M2TUPLE.(M1[1], M2[1])))}(undef, nM1 * nM2) + Ure = zeros(UT, (nU1 * nU2, nM1 * nM2)) + jdx = 0 + for j1 = 1:nM1, j2 = 1:nM2 + jdx += 1 + Mre[jdx] = Vector(M1M2TUPLE.(M1[j1], M2[j2])) + idx = 0 + for i1 = 1:nU1, i2 = 1:nU2 + idx += 1 + Ure[idx, jdx] = U1[i1, j1] * U2[i2, j2] + end + end + + # insert another reduction step + # in my tests this never reduces the size, but I haven't had the time + # to actually prove it isn't needed, so we will keep it for now + Gre = [sum(coco_dot.(Ure[i1, :], Ure[i2, :])) for i1 = 1:size(Ure, 1), i2 = 1:size(Ure, 1)] + Sre = svd(Gre) + rk = rank(Diagonal(Sre.S); rtol=1e-7) + Ure = Sre.U[:, 1:rk]' * Ure + + # now symmetrize w.r.t. permutations + G = _gramian(nn, ll12, Ure, Mre) + S = svd(G) + rk = rank(Diagonal(S.S); rtol=1e-7) + Urpe = S.U[:, 1:rk]' + U = Diagonal(sqrt.(S.S[1:rk])) * Urpe * Ure + + # reconstruct the basis function specifications + rpebs = [PROTOTUPLE.(merge.(nn, ll12, mm12)) for mm12 in Mre] + + return U, rpebs end -function _b2llnn(G::O3O3{L1, M1, L2, M2}, bb) where {L1, M1, L2, M2} - N = length(bb) - @assert all( iszero(b[M]) for b in bb, M in (M1, M2) ) - ll1 = ntuple(i -> bb[i][L1], N) |> SVector - ll2 = ntuple(i -> bb[i][L2], N) |> SVector - nn = ntuple(i -> delete(bb[i], (L1, M1, L2, M2)), N) - ll12 = ntuple(i -> select(bb[i], (L1, L2)), N) - return ll1, ll2, nn, ll12 +function _b2llnn(G::O3O3{L1,M1,L2,M2}, bb) where {L1,M1,L2,M2} + N = length(bb) + @assert all(iszero(b[M]) for b in bb, M in (M1, M2)) + ll1 = ntuple(i -> bb[i][L1], N) |> SVector + ll2 = ntuple(i -> bb[i][L2], N) |> SVector + nn = ntuple(i -> delete(bb[i], (L1, M1, L2, M2)), N) + ll12 = ntuple(i -> select(bb[i], (L1, L2)), N) + return ll1, ll2, nn, ll12 end - diff --git a/src/testing/testing.jl b/src/testing/testing.jl index 8bcf22c..f2fa55b 100644 --- a/src/testing/testing.jl +++ b/src/testing/testing.jl @@ -22,7 +22,7 @@ export print_tf, test_fio, test_transform # # ---------- code for consistency tests # -# test_basis(D::Dict) = ACEfrictionCore.Utils.rpi_basis(; +# test_basis(D::AbstractDict) = ACEfrictionCore.Utils.rpi_basis(; # species = Symbol.(D["species"]), N = D["N"], # maxdeg = D["maxdeg"], # r0 = D["r0"], rcut = D["rcut"], @@ -63,25 +63,25 @@ export print_tf, test_fio, test_transform # import ForwardDiff import ACEfrictionCore: evaluate, inv_transform -function test_transform(T, rrange, ntests = 100) - - rmin, rmax = extrema(rrange) - rr = rmin .+ rand(100) * (rmax-rmin) - xx = evaluate.(Ref(T), rr) - # check syntactic sugar - xx1 = T.(rr) - print_tf(@test xx1 == xx) - # check inversion - rr1 = inv_transform.(Ref(T), xx) - print_tf(@test rr1 ≈ rr) - - # TODO: check that the transform doesn't allocate - @allocated begin - x = 0.0; - for r in rr - x += evaluate(T, r) - end - end +function test_transform(T, rrange, ntests=100) + + rmin, rmax = extrema(rrange) + rr = rmin .+ rand(100) * (rmax - rmin) + xx = evaluate.(Ref(T), rr) + # check syntactic sugar + xx1 = T.(rr) + print_tf(@test xx1 == xx) + # check inversion + rr1 = inv_transform.(Ref(T), xx) + print_tf(@test rr1 ≈ rr) + + # TODO: check that the transform doesn't allocate + @allocated begin + x = 0.0 + for r in rr + x += evaluate(T, r) + end + end end diff --git a/src/transforms/distancetransforms.jl b/src/transforms/distancetransforms.jl index ad2f8cf..9e214d6 100644 --- a/src/transforms/distancetransforms.jl +++ b/src/transforms/distancetransforms.jl @@ -12,17 +12,17 @@ abstract type DistanceTransform end # ----- new transforms implementation import ACEfrictionCore: λ -@deprecate PolyTransform(p, r0) polytransform(p, r0) polytransform(p, r0) = λ("r -> ((1+$r0)/(1+r))^$p") +@deprecate PolyTransform(p, r0) polytransform(p, r0) -@deprecate IdTransform() idtransform() idtransform() = λ("r -> r") +@deprecate IdTransform() idtransform() -@deprecate MorseTransform(lambda, r0) morsetransform(lambda, r0) morsetransform(lambda, r0) = λ("r -> exp(- $lambda * (r / $r0 - 1))") +@deprecate MorseTransform(lambda, r0) morsetransform(lambda, r0) -@deprecate AgnesiTransform(args...) agnesitransform(args...) agnesitransform(r0, p=2, a=(p-1)/(p+1)) = λ("r -> 1 / (1 + $a * (r / $r0)^$p)") +@deprecate AgnesiTransform(args...) agnesitransform(args...) # ------------------------------------------------------ diff --git a/src/transforms/lambdas.jl b/src/transforms/lambdas.jl index 0783468..383714d 100644 --- a/src/transforms/lambdas.jl +++ b/src/transforms/lambdas.jl @@ -26,14 +26,13 @@ lambda(str::String) = λ(str) evaluate(t::Lambda, x) = t.ll.λ(x) -write_dict(t::Lambda) = Dict( - "__id__" => "ACEfrictionCore_Lambda", - "exstr" => t.exstr - ) +write_dict(t::Lambda) = Dict( + "__id__" => "ACEfrictionCore_Lambda", + "exstr" => t.exstr +) -read_dict(::Val{:ACEfrictionCore_Lambda}, D::Dict) = λ(D["exstr"]) +read_dict(::Val{:ACEfrictionCore_Lambda}, D::AbstractDict) = λ(D["exstr"]) import Base: == ==(F1::Lambda, F2::Lambda) = (F1.exstr == F2.exstr) - diff --git a/src/transforms/statetransforms.jl b/src/transforms/statetransforms.jl index e1a4638..f7c2a57 100644 --- a/src/transforms/statetransforms.jl +++ b/src/transforms/statetransforms.jl @@ -1,64 +1,63 @@ using StaticArrays -import ACEfrictionCore -import ACEfrictionCore: evaluate, - write_dict, read_dict, - DState - # frule_evaluate! +import ACEfrictionCore +import ACEfrictionCore: evaluate, + write_dict, read_dict, + DState +# frule_evaluate! -using LinearAlgebra: I, norm +using LinearAlgebra: I, norm -# TODO: +# TODO: # - retire GetNorm -# - polish GetVal, and expand to GetVali -# - polish the frule and rrule implementations +# - polish GetVal, and expand to GetVali +# - polish the frule and rrule implementations -# ------------------ Some different ways to produce an argument +# ------------------ Some different ways to produce an argument -abstract type StateTransform end -abstract type StaticGet <: StateTransform end +abstract type StateTransform end +abstract type StaticGet <: StateTransform end ACEfrictionCore.evaluate(fval::StaticGet, X) = getval(X, fval) valtype(fval::StaticGet, X) = typeof(evaluate(fval, X)) -struct GetVal{VSYM} <: StaticGet end +struct GetVal{VSYM} <: StaticGet end -getval(X, ::GetVal{VSYM}) where {VSYM} = getproperty(X, VSYM) +getval(X, ::GetVal{VSYM}) where {VSYM} = getproperty(X, VSYM) _one(x::Number) = one(x) -_one(x::SVector{3, T}) where {T} = SMatrix{3, 3, T}(I) +_one(x::SVector{3,T}) where {T} = SMatrix{3,3,T}(I) get_symbols(::GetVal{VSYM}) where {VSYM} = (VSYM,) - -# TODO - this is incomplete for now -# struct GetVali{VSYM, IND} <: StaticGet end + +# TODO - this is incomplete for now +# struct GetVali{VSYM, IND} <: StaticGet end # getval(X, ::GetVali{VSYM, IND}) where {VSYM, IND} = getproperty(X, VSYM)[IND] # getval_d(X, ::GetVali{VSYM, IND}) where {VSYM, IND} = __e(getproperty(X, VSYM), Val{IND}()) -struct GetNorm{VSYM} <: StaticGet end +struct GetNorm{VSYM} <: StaticGet end getval(X, ::GetNorm{VSYM}) where {VSYM} = norm(getproperty(X, VSYM)) function getval_d(X, ::GetNorm{VSYM}) where {VSYM} - x = getproperty(X, VSYM) - return DState( NamedTuple{(VSYM,)}( (x/norm(x),) ) ) -end + x = getproperty(X, VSYM) + return DState(NamedTuple{(VSYM,)}((x / norm(x),))) +end get_symbols(::GetNorm{VSYM}) where {VSYM} = (VSYM,) -write_dict(fval::StaticGet) = Dict("__id__" => "ACEfrictionCore_StaticGet", - "expr" => string(typeof(fval)) ) +write_dict(fval::StaticGet) = Dict("__id__" => "ACEfrictionCore_StaticGet", + "expr" => string(typeof(fval))) -read_dict(::Val{:ACEfrictionCore_StaticGet}, D::Dict) = eval( Meta.parse(D["expr"]) )() +read_dict(::Val{:ACEfrictionCore_StaticGet}, D::AbstractDict) = eval(Meta.parse(D["expr"]))() # rrule_evaluate! function removed - derivative functionality has been removed # grad_type_dP function removed - derivative functionality has been removed - diff --git a/src/utils/auxiliary.jl b/src/utils/auxiliary.jl index 7f3dc61..d519a4e 100644 --- a/src/utils/auxiliary.jl +++ b/src/utils/auxiliary.jl @@ -3,21 +3,21 @@ """ random vector on unit sphere, should be uniformly distributed? """ -function rand_sphere(T = Float64) - R = randn(SVector{3, T}) - return R / norm(R) +function rand_sphere(T=Float64) + R = randn(SVector{3,T}) + return R / norm(R) end """ -random rotation matrix; +random rotation matrix; WARN: never checked what the distribution is """ -rand_rot() = (K = (@SMatrix rand(3,3)) .- 0.5; exp(K - K')) +rand_rot() = (K = (@SMatrix rand(3, 3)) .- 0.5; exp(K - K')) """ random reflection, represented as 1, -1 Integer """ -rand_refl() = rand([-1,1]) +rand_refl() = rand([-1, 1]) """ random isometry i.e. element of O(3) @@ -27,22 +27,22 @@ rand_O3() = rand_refl() * rand_rot() -# the following two functions are a little hack to make sure +# the following two functions are a little hack to make sure # that the basis spec is read in the same symbol-order as it is written # (since dicts don't have a specified ordering...) -function _write_dict_1pspec(spec::Vector{NamedTuple{SYMS, NTuple{NSYM, Int}}}) where {SYMS, NSYM} - inds = Vector{Int}[] - for b in spec - vals = [getproperty(b, sym) for sym in SYMS] - push!(inds, vals) - end - return Dict("SYMS" => [ string.(SYMS)... ], "inds" => inds) +function _write_dict_1pspec(spec::Vector{NamedTuple{SYMS,NTuple{NSYM,Int}}}) where {SYMS,NSYM} + inds = Vector{Int}[] + for b in spec + vals = [getproperty(b, sym) for sym in SYMS] + push!(inds, vals) + end + return Dict("SYMS" => [string.(SYMS)...], "inds" => inds) end -function _read_dict_1pspec(D::Dict) - NTPROTO = namedtuple(D["SYMS"]...) - return [ NTPROTO(binds) for binds in D["inds"] ] +function _read_dict_1pspec(D::AbstractDict) + NTPROTO = namedtuple(D["SYMS"]...) + return [NTPROTO(binds) for binds in D["inds"]] end @@ -54,21 +54,21 @@ end """ returns an `SVector{N}` of the form `x * e_I` where `e_I` is the Ith canonical basis vector. """ -@generated function __e(::SVector{N}, ::Val{I}, x::T) where {N, I, T} - code = "SA[" - for i = 1:N - if i == I - code *= "x," - else - code *= "0," - end - end - code *= "]" - quote - $( Meta.parse(code) ) - end +@generated function __e(::SVector{N}, ::Val{I}, x::T) where {N,I,T} + code = "SA[" + for i = 1:N + if i == I + code *= "x," + else + code *= "0," + end + end + code *= "]" + quote + $(Meta.parse(code)) + end end -__e(xx::SVector{N, T}, valI::Val{I}) where {N, T, I} = __e(xx, valI, one(T)) +__e(xx::SVector{N,T}, valI::Val{I}) where {N,T,I} = __e(xx, valI, one(T)) __e(::Number, ::Any, x) = x