Rewrite DNNFriction using Lux

Author: enigne

Project.toml

@@ -7,30 +7,38 @@ version = "0.1.1"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
-Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSOSolvers = "10dff2fc-5484-5881-a0e0-c90441020f8a"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617"
+OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Triangulate = "f7e6ffb2-c36d-4f8f-a77e-16e897189344"
 
 [compat]
 ColorSchemes = "3.25"
 Enzyme = "0.13"
-Flux = "0.16"
-JLD2 = "0.6.2"
+JLD2 = "0.6"
 JSOSolvers = "0.14.1"
+Lux = "1.21"
 MAT = "0.10"
+MLUtils = "0.4.8"
 MadNLP = "0.8"
 Makie = "0.24"
 ManualNLPModels = "0.2.0"
+OnlineStats = "1.7.2"
+Optimisers = "0.4.6"
+Reactant = "0.2.170"
 SparseArrays = "1"
 StatsBase = "0.34"
 Triangulate = "3.0"

src/core/analyses/stressbalanceanalysis.jl

@@ -66,11 +66,8 @@ function UpdateElements(analysis::StressbalanceAnalysis,elements::Vector{Tria},
 		FetchDataToInput(md,inputs,elements,md.friction.C,FrictionCEnum)
 		FetchDataToInput(md,inputs,elements,md.friction.m,FrictionMEnum)
 		FetchDataToInput(md,inputs,elements,md.friction.Cmax,FrictionCmaxEnum)
-	elseif typeof(md.friction) == FluxDNNFriction
-		FetchDataToInput(md,inputs,elements,md.geometry.ssx,SurfaceSlopeXEnum)
-		FetchDataToInput(md,inputs,elements,md.geometry.ssy,SurfaceSlopeYEnum)
-		FetchDataToInput(md,inputs,elements,md.geometry.bsx,BedSlopeXEnum)
-		FetchDataToInput(md,inputs,elements,md.geometry.bsy,BedSlopeYEnum)
+	elseif typeof(md.friction) == DNNFriction
+		FetchDataToInput(md,inputs,elements,md.friction.C,FrictionCEnum)
 	else
 		error("Friction ", typeof(md.friction), " not supported yet")
 	end
@@ -90,11 +87,11 @@ function UpdateParameters(analysis::StressbalanceAnalysis,parameters::Parameters
 		AddParam(parameters, 2, FrictionLawEnum)
 	elseif typeof(md.friction)==SchoofFriction
 		AddParam(parameters, 11, FrictionLawEnum)
-	elseif typeof(md.friction)==FluxDNNFriction
+	elseif typeof(md.friction)==DNNFriction
 		AddParam(parameters, 20, FrictionLawEnum)
-		AddParam(parameters, md.friction.dnnChain, FrictionDNNChainEnum)
-		AddParam(parameters, md.friction.dtx, FrictionDNNdtxEnum)
-		AddParam(parameters, md.friction.dty, FrictionDNNdtyEnum)
+		AddParam(parameters, md.friction.model, FrictionDNNEnum)
+		AddParam(parameters, md.friction.ps, FrictionDNNpsEnum)
+		AddParam(parameters, md.friction.st, FrictionDNNstEnum)
 	else
 		error("Friction ", typeof(md.friction), " not supported yet")
 	end

src/core/friction.jl

@@ -13,7 +13,7 @@ mutable struct CoreBuddFriction <: CoreFriction #{{{
 	rho_water::Float64
 	g::Float64
 end# }}}
-struct  CoreWeertmanFriction <: CoreFriction#{{{
+mutable struct CoreWeertmanFriction <: CoreFriction#{{{
    c_input::Input
 	vx_input::Input
 	vy_input::Input
@@ -31,18 +31,13 @@ mutable struct CoreSchoofFriction <: CoreFriction #{{{
 	rho_water::Float64
 	g::Float64
 end# }}}
-mutable struct CoreFluxDNNFriction <: CoreFriction#{{{
-	dnnChain::Vector{Flux.Chain}
-	dtx::Vector{StatsBase.ZScoreTransform}
-	dty::Vector{StatsBase.ZScoreTransform}
-	xyz_list::Matrix{Float64}
+mutable struct CoreDNNFriction <: CoreFriction#{{{
+	model::AbstractLuxLayer
+	ps
+	st
+   c_input::Input
 	vx_input::Input
 	vy_input::Input
-   b_input::Input
-   H_input::Input
-	rho_ice::Float64
-	rho_water::Float64
-	g::Float64
 end# }}}
 
 function CoreFriction(element::Tria, ::Val{frictionlaw}) where frictionlaw #{{{
@@ -79,23 +74,12 @@ function CoreFriction(element::Tria, ::Val{frictionlaw}) where frictionlaw #{{{
 
 		return CoreSchoofFriction(c_input, vx_input, vy_input, m_input, Cmax_input, H_input, b_input, rho_ice, rho_water, g)
 	elseif frictionlaw==20
-		dnnChain  = FindParam(Vector{Flux.Chain{}}, element, FrictionDNNChainEnum)
-		dtx  = FindParam(Vector{StatsBase.ZScoreTransform{Float64, Vector{Float64}} }, element, FrictionDNNdtxEnum)
-		dty  = FindParam(Vector{StatsBase.ZScoreTransform{Float64, Vector{Float64}} }, element, FrictionDNNdtyEnum)
-		H_input   = GetInput(element, ThicknessEnum)
-		b_input   = GetInput(element, BaseEnum)
-		ssx_input   = GetInput(element, SurfaceSlopeXEnum)
-		ssy_input   = GetInput(element, SurfaceSlopeYEnum)
-		bsx_input   = GetInput(element, BedSlopeXEnum)
-		bsy_input   = GetInput(element, BedSlopeYEnum)
-
-      xyz_list = GetVerticesCoordinates(element.vertices)
-		
-		rho_ice   = FindParam(Float64, element, MaterialsRhoIceEnum)
-		rho_water = FindParam(Float64, element, MaterialsRhoSeawaterEnum)
-		g         = FindParam(Float64, element, ConstantsGEnum)
+		c_input   = GetInput(element, FrictionCEnum)
+		model  = FindParam(AbstractLuxLayer, element, FrictionDNNEnum)
+		ps  = FindParam(NamedTuple, element, FrictionDNNpsEnum)
+		st  = FindParam(NamedTuple, element, FrictionDNNstEnum)
 
-		return CoreFluxDNNFriction(dnnChain,dtx,dty,xyz_list,vx_input,vy_input,b_input,H_input,rho_ice,rho_water,g)
+		return CoreDNNFriction(model,ps,st,c_input,vx_input,vy_input)
 	else
 		error("Friction ",typeof(md.friction)," not supported yet")
 	end
@@ -161,40 +145,18 @@ end#}}}
 	end
 	return alpha2
 end #}}}
-@inline function Alpha2(friction::CoreFluxDNNFriction, gauss::GaussTria, i::Int64)#{{{
-	bed = GetInputValue(friction.b_input, gauss, i)
-	H = GetInputValue(friction.H_input, gauss, i)
+@inline function Alpha2(friction::CoreDNNFriction, gauss::GaussTria, i::Int64)#{{{
 	vx = GetInputValue(friction.vx_input, gauss, i)
 	vy = GetInputValue(friction.vy_input, gauss, i)
-	h = bed + H
 
+	c = GetInputValue(friction.c_input, gauss, i)
 	# Get the velocity
 	vmag = VelMag(friction, gauss, i)
 
-	# velocity gradients
-	dvx = GetInputDerivativeValue(friction.vx_input,friction.xyz_list,gauss,i)
-	dvy = GetInputDerivativeValue(friction.vy_input,friction.xyz_list,gauss,i)
-	vxdx = dvx[1]
-	vxdy = dvx[2]
-	vydx = dvy[1]
-	vydy = dvy[2]
-
-	# Get effective pressure
-	Neff = EffectivePressure(friction, gauss, i)
-
-	# need to change according to the construction of DNN
-	alpha2 = 0.0
-	for i in 1:length(friction.dnnChain)
-		xin = StatsBase.transform(friction.dtx[i], (reshape(vcat(vx, vy, vxdx, vxdy, vydx, vydy, bed, h), 8, :)))
-		pred = StatsBase.reconstruct(friction.dty[i], friction.dnnChain[i](xin))
-		alpha2 += first(pred)
-	end
-	# Average
-	alpha2 = alpha2 / length(friction.dnnChain)
-	if ( (vmag == 0.0) | (alpha2 < 0.0) )
+	if (vmag == 0.0 )
 		alpha2 = 0.0
 	else
-		alpha2 = alpha2 ./ vmag
+		alpha2 = c^2*vmag^2
 	end
 	return alpha2
 end#}}}