Random element assignments

src/utils.jl

@@ -6,7 +6,8 @@ using LinearAlgebra: norm
 
 export rattle!, 
        set_positions, 
-       set_elements
+       set_elements, 
+         randz!
 
 """
 Helper function to convert construct a FlexibleSystem from 
@@ -27,7 +28,10 @@ _set_position(x::Atom, 𝐫) = Atom(atomic_number(x), 𝐫;
                                    velocity = x.velocity,
                                    atomic_mass = x.atomic_mass)
 
-
+_set_element(x::Atom, Z) = Atom(Z, position(x); 
+                                velocity = x.velocity,
+                                atomic_mass = x.atomic_mass)
+
 function set_positions(at::FlexibleSystem, 
                        X::AbstractVector{SVector{3, T}}) where {T}
    @assert length(X) == length(at)                       
@@ -119,6 +123,44 @@ rattle!(sys::FlexibleSystem, r::AbstractFloat) =
       rattle!(sys, r * unit(position(sys)[1][1]))
 
 
+"""
+`randz!(sys::FlexibleSystem, zlist) -> sys`
+
+Randomly assigns elements to the atoms in the system `sys` according to the
+probabilities given in `zlist`.
+`zlist` is an iterable over pairs of the form `id => p` where `id`
+is an atom id (e.g. atomic number or chemical symbol) and `p` 
+a probability. E.g., 
+```
+sys = bulk(:Ti, cubic=true) * 3
+sys = randz!(sys, [ :Ti => 0.2, :O => 0.8 ])
+```
+
+This function was developed mostly for generating testing 
+systems. It may not be suitable for generating random alloys. 
+PRs to improve it are welcome. 
+"""
+function randz!(sys::FlexibleSystem, zlist)
+   ptot = sum( zp[2] for zp in zlist)
+   plist = [ zp[2] / ptot for zp in zlist ]
+   sump = [ sum(plist[1:i]) for i in 1:length(plist) ]
+   sump[end] = 1.001
+   zlist = [ zp[1] for zp in zlist ]   
+
+   for i = 1:length(sys)
+      p = rand() 
+      for j = 1:length(sump)
+         if sump[j] > p
+            sys.particles[i] = _set_element(sys.particles[i], zlist[j])
+            break
+         end
+      end
+   end
+
+   return sys 
+end
+
+
 """
 ```
 union(sys1::FlexibleSystem, sys2::FlexibleSystem)

test/runtests.jl

@@ -4,6 +4,7 @@ using Test
 @testset "AtomsBuilder.jl" begin
     # Write your tests here.
     @testset "bulk" begin include("test_bulk.jl"); end
+    @testset "utils" begin include("test_utils.jl"); end
 end
 
 

test/test_utils.jl

@@ -0,0 +1,12 @@
+
+using AtomsBuilder, Test, AtomsBase, Unitful, Random
+
+##
+
+Random.seed!(1234)
+sys = bulk(:Ti, cubic=true) * 5
+sys = randz!(sys, [ :Ti => 0.5, :O => 0.5 ])
+Z = atomic_number(sys)
+@test count( Z .== 8 ) / length(Z) > 0.45
+@test count( Z .== 22 ) / length(Z) > 0.45
+