Add ViscoElastic trait and extend get_G and get_Kb (#211)

* add ViscoElastic trait and extend get_G and get_Kb * fix testing * increase code coverage * wee fix
JuliaGeodynamics · Aug 15, 2024 · 05c8f1c · 05c8f1c
1 parent beb8e3d
commit 05c8f1c
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 26 deletions.
diff --git a/src/Elasticity/Elasticity.jl b/src/Elasticity/Elasticity.jl
@@ -93,17 +93,8 @@ end
 @inline iselastic(v::AbstractElasticity) = true
 @inline iselastic(v) = false
 
-# for modulus in (:G, :Kb)
-#     fun = Symbol("get_$(string(modulus))")
-#     @eval begin
-#         @inline $(fun)(a::ConstantElasticity) = a.$(modulus).val
-#         @inline $(fun)(a::AbstractMaterialParamsStruct) = isempty(a.Elasticity) ? Inf : $(fun)(a.Elasticity[1])
-#         @inline $(fun)(a::NTuple{N, AbstractMaterialParamsStruct}, phase) where N = nphase($(fun), phase, a)
-#     end
-# end
-
 # extractor methods
-@extractors(ConstantElasticity, :Elasticity)
+# @extractors(ConstantElasticity, :Elasticity)
 
 # Calculation routines
 """

diff --git a/src/GeoParams.jl b/src/GeoParams.jl
@@ -136,9 +136,6 @@ export compute_density,                                # computational routines
 using .MaterialParameters.ConstitutiveRelationships
 export AxialCompression, SimpleShear, Invariant
 
-const get_shearmodulus = get_G
-const get_bulkmodulus = get_Kb
-
 #       Calculation routines
 export dεII_dτII,
     dτII_dεII,
@@ -196,8 +193,6 @@ export dεII_dτII,
     SetConstantElasticity,
     effective_εII,
     iselastic,
-    get_G,
-    get_Kb,
     get_shearmodulus,
     get_bulkmodulus,
 
@@ -340,7 +335,9 @@ export compute_meltfraction,
     SmoothMelting
 
 include("Traits/rheology.jl")
-export islinear, RheologyTrait, LinearRheologyTrait, NonLinearRheologyTrait
+export RheologyTrait
+export islinear, LinearRheologyTrait, NonLinearRheologyTrait
+export isviscoelastic, ElasticRheologyTrait, NonElasticRheologyTrait
 export isplasticity, PlasticRheologyTrait, NonPlasticRheologyTrait
 
 include("Traits/density.jl")
@@ -438,11 +435,22 @@ end
 include("aliases.jl")
 export ntuple_idx
 
-# export DislocationCreep_info,
-#     DiffusionCreep_info,
-#     GrainBoundarySliding_info,
-#     PeierlsCreep_info,
-#     NonLinearPeierlsCreep_info
 
+for modulus in (:G, :Kb)
+    fun = Symbol("get_$(string(modulus))")
+    @eval begin
+        @inline $(fun)(a::ConstantElasticity) = a.$(modulus).val
+        @inline $(fun)(c::CompositeRheology) = $(fun)(isviscoelastic(c), c)
+        @inline $(fun)(::ElasticRheologyTrait, c::CompositeRheology) = mapreduce(x->$(fun)(x), +, c.elements)
+        @inline $(fun)(::AbstractCreepLaw) = 0
+        @inline $(fun)(r::AbstractMaterialParamsStruct) = $(fun)(r.CompositeRheology[1])
+        @inline $(fun)(a::NTuple{N, AbstractMaterialParamsStruct}, phase) where N = nphase($(fun), phase, a)
+    end
+end
+
+export get_G, get_Kb
+
+const get_shearmodulus = get_G
+const get_bulkmodulus  = get_Kb
 
-end # module
+end # module GeoParams
diff --git a/src/Traits/rheology.jl b/src/Traits/rheology.jl
@@ -1,7 +1,9 @@
 abstract type RheologyTrait end
 struct LinearRheologyTrait     <: RheologyTrait end
 struct NonLinearRheologyTrait  <: RheologyTrait end
-struct PlasticRheologyTrait   <: RheologyTrait end
+struct ElasticRheologyTrait    <: RheologyTrait end
+struct NonElasticRheologyTrait <: RheologyTrait end
+struct PlasticRheologyTrait    <: RheologyTrait end
 struct NonPlasticRheologyTrait <: RheologyTrait end
 
 ## LINEAR RHEOLOGY traits
@@ -26,6 +28,28 @@ struct NonPlasticRheologyTrait <: RheologyTrait end
 @inline islinear(r::NTuple{N, Union{AbstractConstitutiveLaw, MaterialParams}}) where N = islinear(islinear(first(r)), islinear(Base.tail(r)))    
 @inline islinear(v::NTuple{1, Union{AbstractConstitutiveLaw, MaterialParams}})         = islinear(v...)
 
+## ELASTICITY RHEOLOGY traits
+
+# traits individual rheologies
+@inline isviscoelastic(::AbstractElasticity)      = ElasticRheologyTrait()
+@inline isviscoelastic(::AbstractConstitutiveLaw) = NonElasticRheologyTrait()
+@inline isviscoelastic(::T) where T               = throw(ArgumentError("$T is an unsupported rheology type"))
+
+# compares two rheologies and return linear trait only and if only both are linear
+@inline isviscoelastic(::RheologyTrait, ::ElasticRheologyTrait) = ElasticRheologyTrait()
+@inline isviscoelastic(::ElasticRheologyTrait, ::RheologyTrait) = ElasticRheologyTrait()
+@inline isviscoelastic(::RheologyTrait, ::RheologyTrait)        = NonElasticRheologyTrait()
+@inline isviscoelastic(v1::AbstractConstitutiveLaw, v2::AbstractConstitutiveLaw) = isviscoelastic(isviscoelastic(v1), isviscoelastic(v2))
+
+# traits for composite rheologies
+@inline isviscoelastic(c::CompositeRheology) = isviscoelastic(c.elements)
+# traits for MaterialParams
+@inline isviscoelastic(r::MaterialParams)    = isviscoelastic(r.CompositeRheology...)
+
+# recursively (pairwise, right-to-left) compare rheology traits of a composite or tuple of material params
+@inline isviscoelastic(r::NTuple{N, Union{AbstractConstitutiveLaw, MaterialParams}}) where N = isviscoelastic(isviscoelastic(first(r)), isviscoelastic(Base.tail(r)))    
+@inline isviscoelastic(v::NTuple{1, Union{AbstractConstitutiveLaw, MaterialParams}})         = isviscoelastic(v...)
+
 ## PLASTIC RHEOLOGY TRAITS
 
 @inline isplasticity(::AbstractPlasticity)      = PlasticRheologyTrait()

diff --git a/test/test_Elasticity.jl b/test/test_Elasticity.jl
@@ -28,10 +28,10 @@ using GeoParams
     @test isvolumetric(a) == true
 
     b = ConstantElasticity()
-    v = SetMaterialParams(; Elasticity=ConstantElasticity())
+    v = SetMaterialParams(; CompositeRheology=CompositeRheology((ConstantElasticity(),)))
     vv = (
-        SetMaterialParams(; Phase=1, Elasticity=ConstantElasticity()),
-        SetMaterialParams(; Phase=2, Elasticity=ConstantElasticity()),
+        SetMaterialParams(; Phase=1, CompositeRheology=CompositeRheology((ConstantElasticity(),))),
+        SetMaterialParams(; Phase=2, CompositeRheology=CompositeRheology((ConstantElasticity(),))),
     )
 
     @test get_G(b) == b.G.val
@@ -42,6 +42,20 @@ using GeoParams
     @test get_Kb(v) == b.Kb.val
     @test get_Kb(vv, 1) == get_Kb(vv, 2) == b.Kb.val # for multiple phases
 
+    # Test get_G and get_Kb with a composite rheology
+    v1 = SetDiffusionCreep(Diffusion.dry_anorthite_Rybacki_2006) # SetPeierlsCreep("Dry Olivine | Goetze and Evans (1979)")
+    v2 = SetDislocationCreep(Dislocation.dry_anorthite_Rybacki_2006)
+    e1 = ConstantElasticity()           # elasticity
+    # CompositeRheologies
+    c = CompositeRheology(v1, v2, e1)   # with elasticity
+
+    r = SetMaterialParams(;
+        CompositeRheology = c
+    )
+
+    @test get_G(c)  == get_G(r)  == r.CompositeRheology[1].elements[3].G.val
+    @test get_Kb(c) == get_Kb(r) == r.CompositeRheology[1].elements[3].Kb.val
+
     # Compute with Floats
     τII = 20e6
     τII_old = 15e6

diff --git a/test/test_Traits.jl b/test/test_Traits.jl
@@ -19,6 +19,9 @@ using Test, GeoParams
         SetMaterialParams(;
             CompositeRheology = CompositeRheology(el, v3),
         ),
+        SetMaterialParams(;
+            CompositeRheology = CompositeRheology(v1, v2),
+        ),
     )
     r2 = (
         SetMaterialParams(;
@@ -32,6 +35,34 @@ using Test, GeoParams
         ),
     )
 
+    # viscoelastic rheology traits
+    # test basic cases
+    for r in (v1, v2, v3, pl)
+        @test isviscoelastic(r) isa NonElasticRheologyTrait
+    end
+    for r in (el,)
+        @test isviscoelastic(r) isa ElasticRheologyTrait
+    end
+    @test_throws ArgumentError isviscoelastic("potato")
+
+    # test composite cases
+    @test isviscoelastic(v1, v2) isa NonElasticRheologyTrait
+    @test isviscoelastic(v1, v3) isa NonElasticRheologyTrait
+    @test isviscoelastic(v3, el) isa ElasticRheologyTrait
+
+    @test isviscoelastic(tuple(v1, v2)) isa NonElasticRheologyTrait
+    @test isviscoelastic(tuple(v1, v3)) isa NonElasticRheologyTrait
+    @test isviscoelastic(tuple(v3, el)) isa ElasticRheologyTrait
+
+    @test isviscoelastic(CompositeRheology(v1, v2, v3)) isa NonElasticRheologyTrait
+    @test isviscoelastic(CompositeRheology(el, v3))     isa ElasticRheologyTrait
+
+    # test MaterialParams cases
+    @test isviscoelastic(r1[1]) isa ElasticRheologyTrait
+    @test isviscoelastic(r1[2]) isa ElasticRheologyTrait
+    @test isviscoelastic(r1[3]) isa ElasticRheologyTrait
+    @test isviscoelastic(r1[4]) isa NonElasticRheologyTrait
+
     ## linear rheology traits
     # test basic cases
     for r in (v1, v2, pl)