Include some special casing for chemical species

src/utils/chemspecies.jl

@@ -1,8 +1,4 @@
-
-
-
-# --------------------------------------------- 
-#  Simple wrapper for chemical element type 
+# Simple wrapper for chemical element type 
 
 import PeriodicTable
 using Unitful
@@ -39,30 +35,29 @@ ChemicalSpecies(:D)
 ```
 """
 struct ChemicalSpecies
-   atomic_number::Int16    # = Z = number of protons
-   nneut::Int16            # number of neutrons
+   atomic_number::Int16  # = Z = number of protons or 0 if unknown
+   n_neutrons::Int16
    info::UInt32
 end
 
 
-function Base.show(io::IO, element::ChemicalSpecies) 
+function Base.show(io::IO, element::ChemicalSpecies)
     print(io, Symbol(element))
 end
 
 Base.Broadcast.broadcastable(s::ChemicalSpecies) = Ref(s)
 
-# better to convert z -> symbol to catch special cases such as D; e.g. 
+# better to convert z -> symbol to catch special cases such as D; e.g.
 # Should ChemicalSpecies(z) == ChemicalSpecies(z,z,0)? For H this is false.
 function ChemicalSpecies(z::Integer; kwargs...)
     ChemicalSpecies(_chem_el_info[z].symbol; kwargs...)
 end
-
 ChemicalSpecies(sym::ChemicalSpecies) = sym
 
-==(a::ChemicalSpecies, sym::Symbol) = 
+==(a::ChemicalSpecies, sym::Symbol) =
         ((a == ChemicalSpecies(sym)) && (Symbol(a) == sym))
 
-# -------- fast access to the periodic table 
+# -------- fast access to the periodic table
 
 if length(PeriodicTable.elements) != maximum(el.number for el in PeriodicTable.elements)
     error("PeriodicTable.elements is not sorted by atomic number")
@@ -72,30 +67,25 @@ if !all(el.number == i for (i, el) in enumerate(PeriodicTable.elements))
     error("PeriodicTable.elements is not sorted by atomic number")
 end
 
-const _chem_el_info = [ 
-      (symbol = Symbol(PeriodicTable.elements[z].symbol), 
-       atomic_mass = PeriodicTable.elements[z].atomic_mass, ) 
-       for z in 1:length(PeriodicTable.elements)
-      ]
+const _chem_el_info = Dict(
+    z => (; symbol = Symbol(PeriodicTable.elements[z].symbol),
+            atomic_mass = PeriodicTable.elements[z].atomic_mass, )
+    for z in 1:length(PeriodicTable.elements)
+)
+_chem_el_info[0] = (; symbol=:X, atomic_mass=0u"u")
+const _sym2z = Dict{Symbol, UInt8}(data.symbol => z for (z, data) in _chem_el_info)
 
-const _sym2z = Dict{Symbol, UInt8}()
-for z in 1:length(_chem_el_info)
-   _sym2z[_chem_el_info[z].symbol] = z
-end
-
-function _nneut_default(z::Integer) 
+function _nneut_default(z::Integer)
     nplusp = round(Int, ustrip(u"u", _chem_el_info[z].atomic_mass))
     return nplusp - z
 end
 
 function ChemicalSpecies(sym::Symbol; n_neutrons = -1, info = 0)
     _islett(c::Char) = 'A' <= uppercase(c) <= 'Z'
 
-    # TODO - special-casing deuterium to make tests pass 
-    #        this should be handled better
-    if sym == :D 
-        return ChemicalSpecies(1, 1, info)
-    end
+    # TODO - special-casing a few things ... This should be handled better
+    sym == :D && return ChemicalSpecies(1, 1, info)
+    sym == :T && return ChemicalSpecies(1, 2, info)
 
     # number of neutrons is explicitly specified
     if n_neutrons != -1
@@ -119,16 +109,15 @@ function ChemicalSpecies(sym::Symbol; n_neutrons = -1, info = 0)
     return ChemicalSpecies(Z, n_neutrons, info)
 end
 
-function Base.Symbol(element::ChemicalSpecies) 
+function Base.Symbol(element::ChemicalSpecies)
     str = "$(_chem_el_info[element.atomic_number].symbol)"
-    if element.nneut != _nneut_default(element.atomic_number)
-        str *= "$(element.atomic_number + element.nneut)"
+    if element.n_neutrons != _nneut_default(element.atomic_number)
+        str *= "$(element.atomic_number + element.n_neutrons)"
     end
 
-    # TODO: again special-casing deuterium; to be fixed. 
-    if str == "H2"
-        return :D
-    end 
+    # TODO: again special-casing deuterium and tritium; to be fixed.
+    str == "H2" && return :D
+    str == "H3" && return :T
 
     return Symbol(str)
 end
@@ -149,10 +138,7 @@ Base.convert(::Type{Symbol}, element::ChemicalSpecies) = Symbol(element)
 
 mass(element::ChemicalSpecies) = _chem_el_info[element.atomic_number].atomic_mass
 
-rich_info(element::ChemicalSpecies) = PeriodicTable.elements[element.atomic_number]
-
-element(element::ChemicalSpecies) = rich_info(element)
-
+element(element::ChemicalSpecies) = PeriodicTable.elements[element.atomic_number]
 
 """The element corresponding to a species/atom (or missing)."""
 element(id::Union{Symbol,Integer}) = PeriodicTable.elements[id]  # Keep for better inlining
@@ -185,10 +171,10 @@ Return the symbols corresponding to the elements of the atoms. Note that
 this may be different than `atomic_symbol` for cases where `atomic_symbol`
 is chosen to be more specific (i.e. designate a special atom).
 """
-element_symbol(sys::AbstractSystem, index) = 
+element_symbol(sys::AbstractSystem, index) =
         element_symbol.(sys[index])
 
-element_symbol(species) = 
+element_symbol(species) =
         Symbol(element(atomic_number(species)).symbol)
 
 

test/species.jl

@@ -1,33 +1,45 @@
-
 using AtomsBase
 using Unitful
 using UnitfulAtomic
 using Test
 
 ##
 
-@testset "ChemicalSpecies" begin 
-
-symbols = [:H, :He, :Li, :Be, :B, :C, :N, :O, :F, :Ne, 
-          :Na, :Mg, :Al, :Si, :P, :S, :Cl, :Ar, :K, :Ca]
-
-# https://thechemicalelements.com/protons-neutrons-electrons-of-elements/
-n_neut =  [0, 2, 4, 5, 6, 6, 7, 8, 10, 10, 
-           12, 12, 14, 14, 16, 16, 18, 22, 20, 20] 
-
-for z = 1:10
-   @test ChemicalSpecies(symbols[z]) == ChemicalSpecies(z) == ChemicalSpecies(z, n_neut[z], 0)
-end
-
-@test ChemicalSpecies(:D) == ChemicalSpecies(1, 1, 0)
-@test ChemicalSpecies(:C13) == ChemicalSpecies(6, 7, 0)
-
-@test atomic_number( UInt(8) ) == 8
-@test atomic_number( Int16(12) ) == 12
-
-@test "$(ChemicalSpecies(:O))" == "$(ChemicalSpecies(8))" == "O"
-@test "$(ChemicalSpecies(8, 8, 0))" == "O"
-@test "$(ChemicalSpecies(:C; n_neutrons=6))" == "C"
-@test "$(ChemicalSpecies(:C; n_neutrons=7))" == "C13"
-
+@testset "ChemicalSpecies" begin
+    @testset "Neutron determination" begin
+        symbols = [:H, :He, :Li, :Be, :B, :C, :N, :O, :F, :Ne,
+                  :Na, :Mg, :Al, :Si, :P, :S, :Cl, :Ar, :K, :Ca]
+
+        # https://thechemicalelements.com/protons-neutrons-electrons-of-elements/
+        n_neut =  [0, 2, 4, 5, 6, 6, 7, 8, 10, 10,
+                   12, 12, 14, 14, 16, 16, 18, 22, 20, 20]
+
+        for z = 1:10
+           @test ChemicalSpecies(symbols[z]) == ChemicalSpecies(z)
+           @test ChemicalSpecies(z) == ChemicalSpecies(z, n_neut[z], 0)
+        end
+
+        @test ChemicalSpecies(:D) == ChemicalSpecies(1, 1, 0)
+        @test ChemicalSpecies(:C13) == ChemicalSpecies(6, 7, 0)
+    end
+
+    @testset "Printing" begin
+        @test "$(ChemicalSpecies(:O))" == "$(ChemicalSpecies(8))" == "O"
+        @test "$(ChemicalSpecies(8, 8, 0))" == "O"
+        @test "$(ChemicalSpecies(:C; n_neutrons=6))" == "C"
+        @test "$(ChemicalSpecies(:C; n_neutrons=7))" == "C13"
+    end
+
+    @testset "Special cases" begin
+        # Test a few special cases that come up in the wild
+        x = ChemicalSpecies(0)
+        @test x.n_neutrons == 0
+        @test atomic_number(x) == 0
+        @test mass(x) == 0u"u"
+
+        @test x = ChemicalSpecies(:X)
+    end
+
+    @test atomic_number( UInt(8) ) == 8
+    @test atomic_number( Int16(12) ) == 12
 end