Skip to content

Commit

Permalink
Merge branch 'master' into gpu
Browse files Browse the repository at this point in the history
  • Loading branch information
@albop
albop authored Oct 22, 2024
2 parents c935029 + 6e4e262 commit cf84bba
Show file tree
Hide file tree
Showing 4 changed files with 220 additions and 1 deletion.
49 changes: 49 additions & 0 deletions misc/test_truncdiff.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
using Dolo: splines


ranges = ((0.0, 1.0, 10), (0.0, 1.0, 20))
methods(splines.interp)

using StaticArrays

x0 = SVector( 10,20)


values = rand(typeof(x0), 10,20)

values2 = rand(typeof(x0), 12,22)

using ForwardDiff


fun = u->splines.interp(ranges, values, u)

ForwardDiff.jacobian(
fun,
x0
)

a = tuple((e[1] for e in ranges)...)
b = tuple((e[2] for e in ranges)...)
orders = tuple((e[3] for e in ranges)...)

fun2 = u->splines.eval_UC_spline(a,b,orders, values2, u)
fun2(x0)




ForwardDiff.jacobian(
fun2,
x0
)

# import ForwardDiff: unsafe_trunc

# function ForwardDiff.unsafe_trunc(TI, number::ForwardDiff.Dual{Tag, Tf}) where Tag where Tf<:Union{Float32, Float64}
# ForwardDiff.Dual{Tag}(
# unsafe_trunc(TI, ForwardDiff.value(number)),
# 0*ForwardDiff.partials(number)
# )

# end
11 changes: 11 additions & 0 deletions src/funs.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,10 @@ end

function (f::DFun{A,B,I,vars})(i::Int, j::Int) where A where B<:GArray{G,V} where V where I where G<:PGrid{G1,G2} where G1<:SGrid where G2<:CGrid where vars where d2 where U
f.values[i,j]

function (f::DFun{A,B,I,vars})(x::QP) where A where B<:GArray{G,V} where V where I where G<:PGrid{G1,G2} where G1<:SGrid where G2<:CGrid where vars
f(x.loc...)
# f(x.loc)
end

function (f::DFun{A,B,I,vars})(x::Tuple) where A where B<:GArray{G,V} where V where I where G<:PGrid{G1,G2} where G1<:SGrid where G2<:CGrid where vars
Expand All @@ -137,6 +141,13 @@ function (f::DFun{A,B,I,vars})(loc::Tuple{Tuple{Int64}, SVector{d2, U}}) where
f.itp[loc[1][1]](x)
end

function (f::DFun{A,B,I,vars})(loc::Tuple{Int64, Int64}) where A where B<:GArray{G,V} where V where I where G<:PGrid{G1,G2} where G1<:SGrid where G2<:CGrid where vars where d2 where U
# TODO: not beautiful
i,j = loc
x = f.values.grid.grids[2][j]
f.itp[loc[1][1]](x)
end

# CGrid

function (f::DFun{A,B,I,vars})(x::QP) where A where B<:GArray{G,V} where V where I where G<:CGrid where vars
Expand Down
157 changes: 157 additions & 0 deletions src/splines/csplines.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,166 @@ end
return fun.args[2].args[2]
end


@generated function eval_UC_spline_(a, b, orders, C::AbstractArray{T, d}, S::SVector{d,U}) where d where T where U
fun = create_function(d,"natural"; vectorize=false)
return fun.args[2].args[2]
end

eval_UC_spline(a, b, orders, C::AbstractArray{T, d}, S::SVector{d,U}) where d where U where T = eval_UC_spline_(a, b, orders, C, S)


function cPhi(λ)

λ0 = 1
λ1 = λ
λ2 = λ*λ1
λ3 = λ*λ2

if λ1<0
Phi_4 = -0.5 * λ1 + 0.16666666666666666
Phi_3 = 0.0 * λ1 + 0.6666666666666666
Phi_2 = 0.5 * λ1 + 0.16666666666666666
Phi_1 = 0.0 * λ1 + 0.0
elseif λ1>1
Phi_4 = (3 * -0.16666666666666666 + 2 * 0.5 + -0.5) * (1 - 1) + (-0.16666666666666666 + 0.5 + -0.5 + 0.16666666666666666)
Phi_3 = (3 * 0.5 + 2 * -1.0 + 0.0) * (λ1 - 1) + (0.5 + -1.0 + 0.0 + 0.6666666666666666)
Phi_2 = (3 * -0.5 + 2 * 0.5 + 0.5) * (λ1 - 1) + (-0.5 + 0.5 + 0.5 + 0.16666666666666666)
Phi_1 = (3 * 0.16666666666666666 + 2 * 0.0 + 0.0) * (λ1 - 1) + (0.16666666666666666 + 0.0 + 0.0 + 0.0)
else
Phi_1 = -0.16666666666666666 * λ3 + 0.5 * λ2 + -0.5 * λ1 + 0.16666666666666666 * λ0
Phi_2 = 0.5 * λ3 + -1.0 * λ2 + 0.0 * λ1 + 0.6666666666666666 * λ0
Phi_3 = -0.5 * λ3 + 0.5 * λ2 + 0.5 * λ1 + 0.16666666666666666 * λ0
Phi_4 = 0.16666666666666666 * λ3 + 0.0 * λ2 + 0.0 * λ1 + 0.0 * λ0
end


return SVector(Phi_4, Phi_3, Phi_2, Phi_1)
end


function mextract(c::AbstractArray{T,d}, inds::NTuple{d,<:Int}, ::Val{n}) where T where d where n
ranges = tuple((i:i+n-1 for i in inds)...)
return @inbounds SArray{NTuple{d,n},T,d,n^d}(view(c, ranges...))
end

@inline function reduce_tensors( v::SArray, Φ::NTuple{d,SVector{n,U}}) where d where n where U
xx = SVector( (prod(e) for e in Iterators.product(Φ...))...)
return sum( xx .* v[:])
end

function reduce_tensors(v::SArray, Φ::Tuple{SVector{4,U}}) where U
@inbounds v[1]*Φ[1][1] + v[2]*Φ[1][2] + v[3]*Φ[1][3] + v[4]*Φ[1][4]
end


function reduce_tensors(v::SArray, Φ::Tuple{SVector{4,U},SVector{4,U}}) where U
@inbounds v[1,1]*Φ[1][1]*Φ[2][1] +
v[1,2]*Φ[1][1]*Φ[2][2] +
v[1,3]*Φ[1][1]*Φ[2][3] +
v[1,4]*Φ[1][1]*Φ[2][4] +
v[2,1]*Φ[1][2]*Φ[2][1] +
v[2,2]*Φ[1][2]*Φ[2][2] +
v[2,3]*Φ[1][2]*Φ[2][3] +
v[2,4]*Φ[1][2]*Φ[2][4] +
v[3,1]*Φ[1][3]*Φ[2][1] +
v[3,2]*Φ[1][3]*Φ[2][2] +
v[3,3]*Φ[1][3]*Φ[2][3] +
v[3,4]*Φ[1][3]*Φ[2][4] +
v[4,1]*Φ[1][4]*Φ[2][1] +
v[4,2]*Φ[1][4]*Φ[2][2] +
v[4,3]*Φ[1][4]*Φ[2][3] +
v[4,4]*Φ[1][4]*Φ[2][4]
end

function reduce_tensors(v::SArray, Φ::Tuple{SVector{4,U},SVector{4,U},SVector{4,U}}) where U
a = @inbounds v[1,1,1]*Φ[1][1]*Φ[2][1]*Φ[3][1] +
v[1,1,2]*Φ[1][1]*Φ[2][1]*Φ[3][2] +
v[1,1,3]*Φ[1][1]*Φ[2][1]*Φ[3][3] +
v[1,1,4]*Φ[1][1]*Φ[2][1]*Φ[3][4] +
v[1,2,1]*Φ[1][1]*Φ[2][2]*Φ[3][1] +
v[1,2,2]*Φ[1][1]*Φ[2][2]*Φ[3][2] +
v[1,2,3]*Φ[1][1]*Φ[2][2]*Φ[3][3] +
v[1,2,4]*Φ[1][1]*Φ[2][2]*Φ[3][4] +
v[1,3,1]*Φ[1][1]*Φ[2][3]*Φ[3][1] +
v[1,3,2]*Φ[1][1]*Φ[2][3]*Φ[3][2] +
v[1,3,3]*Φ[1][1]*Φ[2][3]*Φ[3][3] +
v[1,3,4]*Φ[1][1]*Φ[2][3]*Φ[3][4] +
v[1,4,1]*Φ[1][1]*Φ[2][4]*Φ[3][1] +
v[1,4,2]*Φ[1][1]*Φ[2][4]*Φ[3][2] +
v[1,4,3]*Φ[1][1]*Φ[2][4]*Φ[3][3] +
v[1,4,4]*Φ[1][1]*Φ[2][4]*Φ[3][4]
b = @inbounds v[2,1,1]*Φ[1][2]*Φ[2][1]*Φ[3][1] +
v[2,1,2]*Φ[1][2]*Φ[2][1]*Φ[3][2] +
v[2,1,3]*Φ[1][2]*Φ[2][1]*Φ[3][3] +
v[2,1,4]*Φ[1][2]*Φ[2][1]*Φ[3][4] +
v[2,2,1]*Φ[1][2]*Φ[2][2]*Φ[3][1] +
v[2,2,2]*Φ[1][2]*Φ[2][2]*Φ[3][2] +
v[2,2,3]*Φ[1][2]*Φ[2][2]*Φ[3][3] +
v[2,2,4]*Φ[1][2]*Φ[2][2]*Φ[3][4] +
v[2,3,1]*Φ[1][2]*Φ[2][3]*Φ[3][1] +
v[2,3,2]*Φ[1][2]*Φ[2][3]*Φ[3][2] +
v[2,3,3]*Φ[1][2]*Φ[2][3]*Φ[3][3] +
v[2,3,4]*Φ[1][2]*Φ[2][3]*Φ[3][4] +
v[2,4,1]*Φ[1][2]*Φ[2][4]*Φ[3][1] +
v[2,4,2]*Φ[1][2]*Φ[2][4]*Φ[3][2] +
v[2,4,3]*Φ[1][2]*Φ[2][4]*Φ[3][3] +
v[2,4,4]*Φ[1][2]*Φ[2][4]*Φ[3][4]
c = @inbounds v[3,1,1]*Φ[1][3]*Φ[2][1]*Φ[3][1] +
v[3,1,2]*Φ[1][3]*Φ[2][1]*Φ[3][2] +
v[3,1,3]*Φ[1][3]*Φ[2][1]*Φ[3][3] +
v[3,1,4]*Φ[1][3]*Φ[2][1]*Φ[3][4] +
v[3,2,1]*Φ[1][3]*Φ[2][2]*Φ[3][1] +
v[3,2,2]*Φ[1][3]*Φ[2][2]*Φ[3][2] +
v[3,2,3]*Φ[1][3]*Φ[2][2]*Φ[3][3] +
v[3,2,4]*Φ[1][3]*Φ[2][2]*Φ[3][4] +
v[3,3,1]*Φ[1][3]*Φ[2][3]*Φ[3][1] +
v[3,3,2]*Φ[1][3]*Φ[2][3]*Φ[3][2] +
v[3,3,3]*Φ[1][3]*Φ[2][3]*Φ[3][3] +
v[3,3,4]*Φ[1][3]*Φ[2][3]*Φ[3][4] +
v[3,4,1]*Φ[1][3]*Φ[2][4]*Φ[3][1] +
v[3,4,2]*Φ[1][3]*Φ[2][4]*Φ[3][2] +
v[3,4,3]*Φ[1][3]*Φ[2][4]*Φ[3][3] +
v[3,4,4]*Φ[1][3]*Φ[2][4]*Φ[3][4]
d = @inbounds v[4,1,1]*Φ[1][4]*Φ[2][1]*Φ[3][1] +
v[4,1,2]*Φ[1][4]*Φ[2][1]*Φ[3][2] +
v[4,1,3]*Φ[1][4]*Φ[2][1]*Φ[3][3] +
v[4,1,4]*Φ[1][4]*Φ[2][1]*Φ[3][4] +
v[4,2,1]*Φ[1][4]*Φ[2][2]*Φ[3][1] +
v[4,2,2]*Φ[1][4]*Φ[2][2]*Φ[3][2] +
v[4,2,3]*Φ[1][4]*Φ[2][2]*Φ[3][3] +
v[4,2,4]*Φ[1][4]*Φ[2][2]*Φ[3][4] +
v[4,3,1]*Φ[1][4]*Φ[2][3]*Φ[3][1] +
v[4,3,2]*Φ[1][4]*Φ[2][3]*Φ[3][2] +
v[4,3,3]*Φ[1][4]*Φ[2][3]*Φ[3][3] +
v[4,3,4]*Φ[1][4]*Φ[2][3]*Φ[3][4] +
v[4,4,1]*Φ[1][4]*Φ[2][4]*Φ[3][1] +
v[4,4,2]*Φ[1][4]*Φ[2][4]*Φ[3][2] +
v[4,4,3]*Φ[1][4]*Φ[2][4]*Φ[3][3] +
v[4,4,4]*Φ[1][4]*Φ[2][4]*Φ[3][4]
a+b+c+d
end

function eval_spline( ranges, C::AbstractArray{T,d}, x::SVector{d,U}) where d where U where T

a = SVector( (e[1] for e in ranges)... )
b = SVector( (e[2] for e in ranges)... )
n = SVector( (e[3] for e in ranges)... )

δ = (b.-a)./(n.-1)
i = div.( (x.-a), δ)
i = max.(min.(i, n.-2), 0)
λ = (x.-(a .+ δ.*i))./δ

# i_ = floor.(Int, i) .+ 1
i_ = unsafe_trunc.(Int, i) .+ 1

Φ = tuple( (cPhi(e) for e in λ)... )

M = mextract(C, tuple( (ii for ii in i_)...), Val(4))

# return sum(M) + sum(Φ[1])+ sum(Φ[2]) + sum(Φ[3])
# reduce_tensors(M,Φ)
@inline reduce_tensors(M,Φ)

end
4 changes: 3 additions & 1 deletion src/splines/macros.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,9 @@ function create_Phi(d, extrap, diff; Tf=Float64)
block = []
rhs_1 = U("tp", i,1)
rhs_2 = U("tp", i,2)
rhs_4 = U("tp", i,4)
rhs_3 = U("tp", i,3)
rhs_4 = U("tp", i,4)

if extrap == "none"
for j=1:4
eq = :($(U("Phi_",i,j)) = ($(Ad[j,1])*$rhs_1 + $(Ad[j,2])*$rhs_2 + $(Ad[j,3])*$rhs_3 + $(Ad[j,4])*$rhs_4) )
Expand Down Expand Up @@ -221,6 +222,7 @@ function create_function(d,extrap="natural"; vectorize=true, Tf=Float64)
$(create_parameters(d; Tf=Tf)...)
N = size(S,1)
# @fastmath @inbounds @simd( # doesn't seem to make any difference

for n=1:N
$(create_local_parameters(d;Tf=Tf)...)
$(create_Phi(d,extrap,false;Tf=Tf)...)
Expand Down

0 comments on commit cf84bba

Please sign in to comment.