fix: made models continuous

Author: hack-hard

src/layers.jl

@@ -141,8 +141,8 @@ function symetrise(val::StructArray{PreprocessedData{T}};
     vcat(dot,
         r_1 .+ r_2,
         abs.(r_1 .- r_2),
-        d_1 .+ d_2, abs.(d_1 .- d_2)) .*
-    cut.(cutoff_radius, r_1) .* cut.(cutoff_radius, r_2)
+        d_1 .+ d_2, abs.(d_1 .- d_2),
+        cut.(cutoff_radius, r_1) .* cut.(cutoff_radius, r_2))
 end
 scale_factor(x) = x[end:end, :]
 

src/models.jl

@@ -19,9 +19,9 @@ function evaluate_if_atoms_in_neighboord(layer, arg::AbstractArray, ps, st; zero
 end
 
 function general_angular_dense(main_chain, secondary_chain; name::String,
-        van_der_waals_channel = false, on_gpu = true, cutoff_radius::Float32 = 3.0f0)
+    van_der_waals_channel=false, on_gpu=true, cutoff_radius::Float32=3.0f0)
     main_chain = DeepSet(Chain(
-        symetrise(; cutoff_radius, device = on_gpu ? gpu_device() : identity),
+        symetrise(; cutoff_radius, device=on_gpu ? gpu_device() : identity),
         main_chain
     ))
     function add_van_der_waals_channel(main_chain)
@@ -41,44 +41,53 @@ end
 	`tiny_angular_dense` is a function that generate a lux model.
 
 """
-function tiny_angular_dense(; van_der_waals_channel = false, kargs...)
+function tiny_angular_dense(; van_der_waals_channel=false, kargs...)
     general_angular_dense(
-        Chain(Dense(5 => 7, elu),
-            Dense(7 => 4, elu)),
+        Parallel(.*,
+            Chain(Dense(6 => 7, elu),
+                Dense(7 => 4, elu)),
+            Lux.WrappedFunction(scale_factor)
+        ),
         Chain(
             BatchNorm(4 + van_der_waals_channel),
             Dense(4 + van_der_waals_channel => 6, elu),
             Dense(6 => 1, sigmoid_fast));
-        name = "tiny_angular_dense_" *
-               (van_der_waals_channel ? "v" : ""),
+        name="tiny_angular_dense_" *
+             (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel, kargs...)
 end
 
-function light_angular_dense(; van_der_waals_channel = false, kargs...)
+function light_angular_dense(; van_der_waals_channel=false, kargs...)
     general_angular_dense(
-        Chain(Dense(5 => 10, elu),
-            Dense(10 => 5, elu)),
+        Parallel(.*,
+            Chain(Dense(6 => 10, elu),
+                Dense(10 => 5, elu)),
+            Lux.WrappedFunction(scale_factor)
+        ),
         Chain(
             BatchNorm(5 + van_der_waals_channel),
             Dense(5 + van_der_waals_channel => 10, elu),
             Dense(10 => 1, sigmoid_fast));
-        name = "light_angular_dense_" *
-               (van_der_waals_channel ? "v" : ""),
+        name="light_angular_dense_" *
+             (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel, kargs...)
 end
 
 function medium_angular_dense(;
-        van_der_waals_channel = false, kargs...)
-    general_angular_dense(Chain(
-            Dense(5 => 15, elu),
-            Dense(15 => 10, elu)),
+    van_der_waals_channel=false, kargs...)
+    general_angular_dense(
+        Parallel(.*,
+            Chain(Dense(6 => 15, elu),
+                Dense(15 => 5, elu)),
+            Lux.WrappedFunction(scale_factor)
+        ),
         Chain(
             BatchNorm(10 + van_der_waals_channel),
-            Dense(10 + van_der_waals_channel => 5; use_bias = false),
+            Dense(10 + van_der_waals_channel => 5; use_bias=false),
             Dense(5 => 10, elu),
             Dense(10 => 1, sigmoid_fast));
-        name = "medium_angular_dense_" *
-               (van_der_waals_channel ? "v" : ""),
+        name="medium_angular_dense_" *
+             (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel,
         kargs...)
 end