Skip to content

Commit

Permalink
check AoS v SoA
Browse files Browse the repository at this point in the history
  • Loading branch information
mattsignorelli committed Sep 10, 2024
1 parent af2c8d8 commit bf6a1ee
Show file tree
Hide file tree
Showing 4 changed files with 232 additions and 6 deletions.
150 changes: 149 additions & 1 deletion Manifest.toml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

julia_version = "1.10.4"
manifest_format = "2.0"
project_hash = "266de5c5e4a1f3d9c635464e38a9db2a7626d65d"
project_hash = "ff2db58e7d57852b3d706870ee21f2a32b4e69fb"

[[deps.AcceleratorLattice]]
deps = ["Accessors", "EnumX", "InteractiveUtils", "LinearAlgebra", "PyFormattedStrings", "ReferenceFrameRotations", "Test"]
Expand Down Expand Up @@ -33,6 +33,12 @@ version = "0.1.37"
StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"

[[deps.AliasTables]]
deps = ["PtrArrays", "Random"]
git-tree-sha1 = "9876e1e164b144ca45e9e3198d0b689cadfed9ff"
uuid = "66dad0bd-aa9a-41b7-9441-69ab47430ed8"
version = "1.1.3"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
version = "1.1.1"
Expand Down Expand Up @@ -62,6 +68,16 @@ git-tree-sha1 = "bce6804e5e6044c6daab27bb533d1295e4a2e759"
uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
version = "0.7.6"

[[deps.Compat]]
deps = ["TOML", "UUIDs"]
git-tree-sha1 = "8ae8d32e09f0dcf42a36b90d4e17f5dd2e4c4215"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "4.16.0"
weakdeps = ["Dates", "LinearAlgebra"]

[deps.Compat.extensions]
CompatLinearAlgebraExt = "LinearAlgebra"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
Expand Down Expand Up @@ -107,6 +123,12 @@ git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.16.0"

[[deps.DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.20"

[[deps.DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
Expand All @@ -116,6 +138,22 @@ version = "1.0.0"
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.Distributions]]
deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"]
git-tree-sha1 = "e6c693a0e4394f8fda0e51a5bdf5aef26f8235e9"
uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
version = "0.25.111"

[deps.Distributions.extensions]
DistributionsChainRulesCoreExt = "ChainRulesCore"
DistributionsDensityInterfaceExt = "DensityInterface"
DistributionsTestExt = "Test"

[deps.Distributions.weakdeps]
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d"
Expand All @@ -141,6 +179,18 @@ version = "0.1.10"
[[deps.FileWatching]]
uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"

[[deps.FillArrays]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "6a70198746448456524cb442b8af316927ff3e1a"
uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
version = "1.13.0"
weakdeps = ["PDMats", "SparseArrays", "Statistics"]

[deps.FillArrays.extensions]
FillArraysPDMatsExt = "PDMats"
FillArraysSparseArraysExt = "SparseArrays"
FillArraysStatisticsExt = "Statistics"

[[deps.GTPSA]]
deps = ["GTPSA_jll", "LinearAlgebra", "PrettyTables", "Printf", "SpecialFunctions"]
git-tree-sha1 = "3c0daebf55c4a861accf79106b55843fffb636e0"
Expand All @@ -159,6 +209,12 @@ git-tree-sha1 = "d1d712be3164d61d1fb98e7ce9bcbc6cc06b45ed"
uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
version = "1.10.8"

[[deps.HypergeometricFunctions]]
deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"]
git-tree-sha1 = "7c4195be1649ae622304031ed46a2f4df989f1eb"
uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a"
version = "0.3.24"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
Expand Down Expand Up @@ -277,6 +333,12 @@ deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
version = "2.28.2+1"

[[deps.Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "ec4f7fbeab05d7747bdf98eb74d130a2a2ed298d"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.2.0"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

Expand Down Expand Up @@ -327,6 +389,12 @@ git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.6.3"

[[deps.PDMats]]
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
git-tree-sha1 = "949347156c25054de2db3b166c52ac4728cbad65"
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
version = "0.11.31"

[[deps.Parsers]]
deps = ["Dates", "PrecompileTools", "UUIDs"]
git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821"
Expand Down Expand Up @@ -360,12 +428,29 @@ version = "2.3.2"
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.PtrArrays]]
git-tree-sha1 = "77a42d78b6a92df47ab37e177b2deac405e1c88f"
uuid = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d"
version = "1.2.1"

[[deps.PyFormattedStrings]]
deps = ["Printf", "SnoopPrecompile"]
git-tree-sha1 = "52d272f8045d6787cb5f45b6273dcf4634acc2d4"
uuid = "5f89f4a4-a228-4886-b223-c468a82ed5b9"
version = "0.1.11"

[[deps.QuadGK]]
deps = ["DataStructures", "LinearAlgebra"]
git-tree-sha1 = "1d587203cf851a51bf1ea31ad7ff89eff8d625ea"
uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
version = "2.11.0"

[deps.QuadGK.extensions]
QuadGKEnzymeExt = "Enzyme"

[deps.QuadGK.weakdeps]
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
Expand All @@ -385,6 +470,18 @@ git-tree-sha1 = "c1fadb2e94fa53c8bb8b937118e87d388012eaf6"
uuid = "74f56ac7-18b3-5285-802d-d4bd4f104033"
version = "3.0.2"

[[deps.Rmath]]
deps = ["Random", "Rmath_jll"]
git-tree-sha1 = "f65dcb5fa46aee0cf9ed6274ccbd597adc49aa7b"
uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa"
version = "0.7.1"

[[deps.Rmath_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl"]
git-tree-sha1 = "e60724fd3beea548353984dc61c943ecddb0e29a"
uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
version = "0.4.3+0"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
version = "0.7.0"
Expand All @@ -406,6 +503,17 @@ version = "1.0.3"
[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.2.1"

[[deps.SparseArrays]]
deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
version = "1.10.0"

[[deps.SpecialFunctions]]
deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14"
Expand Down Expand Up @@ -437,12 +545,52 @@ git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682"
uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
version = "1.4.3"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
version = "1.10.0"

[[deps.StatsAPI]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed"
uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
version = "1.7.0"

[[deps.StatsBase]]
deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
git-tree-sha1 = "5cf7606d6cef84b543b483848d4ae08ad9832b21"
uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
version = "0.34.3"

[[deps.StatsFuns]]
deps = ["HypergeometricFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"]
git-tree-sha1 = "cef0472124fab0695b58ca35a77c6fb942fdab8a"
uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
version = "1.3.1"

[deps.StatsFuns.extensions]
StatsFunsChainRulesCoreExt = "ChainRulesCore"
StatsFunsInverseFunctionsExt = "InverseFunctions"

[deps.StatsFuns.weakdeps]
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"

[[deps.StringManipulation]]
deps = ["PrecompileTools"]
git-tree-sha1 = "a04cabe79c5f01f4d723cc6704070ada0b9d46d5"
uuid = "892a3eda-7b42-436c-8928-eab12a02cf0e"
version = "0.3.4"

[[deps.SuiteSparse]]
deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"

[[deps.SuiteSparse_jll]]
deps = ["Artifacts", "Libdl", "libblastrampoline_jll"]
uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c"
version = "7.2.1+1"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
Expand Down
1 change: 1 addition & 0 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ version = "1.0.0-DEV"
[deps]
AcceleratorLattice = "118f53b2-f4ec-480a-a460-1a04f660b17d"
AtomicAndPhysicalConstants = "5c0d271c-5419-4163-b387-496237733d8b"
Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
ReferenceFrameRotations = "74f56ac7-18b3-5285-802d-d4bd4f104033"
Expand Down
68 changes: 65 additions & 3 deletions src/BeamTracking.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,21 @@ module BeamTracking
using AcceleratorLattice,
GTPSA,
ReferenceFrameRotations,
StaticArrays
StaticArrays,
Distributions

# Temporary until AtomicAndPhysicalConstants is cleaned up ----
import AtomicAndPhysicalConstants: AtomicAndPhysicalConstants
const Species = AtomicAndPhysicalConstants.Particle
# -------------------------------------------------------------

export Beam,
export Beam,
Coords,
Coord,
Particle,
Symplectic,
Paraxial
Paraxial,
Species

# SoA ----------------------------------
struct Coords{T} <: FieldVector{6, T}
Expand All @@ -23,10 +28,39 @@ struct Coords{T} <: FieldVector{6, T}
pz::Vector{T}
end

function Coords(
n::Integer;
d_x::Distribution=Normal(0,0), d_px::Distribution=Normal(0,0),
d_y::Distribution=Normal(0,0), d_py::Distribution=Normal(0,0),
d_z::Distribution=Normal(0,0), d_pz::Distribution=Normal(0,0)
)
x = rand(d_x , n)
px = rand(d_px, n)
y = rand(d_y , n)
py = rand(d_py, n)
z = rand(d_z , n)
pz = rand(d_pz, n)

return Coords(x, px, y, py, z, pz)
end

struct Beam{T} # Must agree exactly with `Particle`!
species::Species
z::Coords{T}
end

function Beam(
n::Integer; species::Species=Species("electron"),
d_x::Distribution=Normal(0,0), d_px::Distribution=Normal(0,0),
d_y::Distribution=Normal(0,0), d_py::Distribution=Normal(0,0),
d_z::Distribution=Normal(0,0), d_pz::Distribution=Normal(0,0)
)

coords = Coords(n; d_x=d_x, d_px=d_px, d_y=d_y, d_py=d_py, d_z=d_z, d_pz=d_pz)

return Beam(species, coords)
end

# --------------------------------------

# AoS ----------------------------------
Expand All @@ -39,10 +73,38 @@ struct Coord{T} <: FieldVector{6, T}
pz::T
end

function Coord(;
d_x::Distribution=Normal(0,0), d_px::Distribution=Normal(0,0),
d_y::Distribution=Normal(0,0), d_py::Distribution=Normal(0,0),
d_z::Distribution=Normal(0,0), d_pz::Distribution=Normal(0,0)
)
x = rand(d_x )
px = rand(d_px)
y = rand(d_y )
py = rand(d_py)
z = rand(d_z )
pz = rand(d_pz)

return Coord(x, px, y, py, z, pz)
end

struct Particle{T} # Must agree exactly with `Beam`!
species::Species
z::Coord{T}
end

function Particle(;
species::Species=Species("electron"),
d_x::Distribution=Normal(0,0), d_px::Distribution=Normal(0,0),
d_y::Distribution=Normal(0,0), d_py::Distribution=Normal(0,0),
d_z::Distribution=Normal(0,0), d_pz::Distribution=Normal(0,0)
)

coord = Coord(d_x=d_x, d_px=d_px, d_y=d_y, d_py=d_py, d_z=d_z, d_pz=d_pz)

return Particle(species, coord)
end

# ---------------------------------------


Expand Down
19 changes: 17 additions & 2 deletions src/paraxial/Paraxial.jl
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,9 @@ where `statef` and `state0` are `Vector{Particle{T}}` one can do
- `statef` -- Output state after tracking through, may be either a `Particle` or a `Beam`
- `state0` -- Input state before tracking through, may be either a `Particle` or a `Beam`
"""
function track!(ele::Drift, statef::T, state0::T) where {T <: Union{Beam,Particle}}
function track!(L, statef::Beam{T}, state0::Beam{T}) where {T}
@assert !(statef === state0) "Aliasing statef === state0 not allowed!"
L = ele.L
#L = 5
z0 = state0.z
zf = statef.z

Expand All @@ -36,4 +36,19 @@ function track!(ele::Drift, statef::T, state0::T) where {T <: Union{Beam,Particl
return statef
end

function track_one(L, state0::Particle)
#L = ele.L
z0 = state0.z

@FastGTPSA begin
x = z0[1]+z0[2]*L/(1.0+z0[6])
px = z0[2]
y = z0[3]+z0[4]*L/(1.0+z0[6])
py = z0[4]
z = z0[5]-L*((z0[2]^2)+(z0[4]^2))/(1.0+z0[6])^2/2.0
pz = z0[6]
end
return Particle(state0.species, Coord(x,px,y,py,z,pz))
end

end

0 comments on commit bf6a1ee

Please sign in to comment.