Phase Space Layout (QEDjl-project#70)

This PR adds the phase space layout interface and a convenient
user-faced API for working with coordinate maps for given processes,
models, and phase space layouts.

# TODO: 

- [x] clean up (remove the TODOs in the files)
- [x] unify the calls of coordinate maps (are the different
implementations for in- and out-PSL vital?)
- [x] add validity checks to the interface and construction of the
coordinate maps
- [x] fix the docs
- [x] move `patch_QEDbase.jl` to QEDbase

Project.toml

@@ -9,13 +9,14 @@ QEDbase = "10e22c08-3ccb-4172-bfcf-7d7aa3d04d93"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
-DocStringExtensions = "0.9"
-QEDbase = "0.3.0"
+DocStringExtensions = "^0.9"
 Reexport = "1.2"
 SimpleTraits = "0.9"
 StaticArrays = "1.9"
+StatsBase = "0.34.3"
 julia = "1.10"
 
 [extras]

docs/make.jl

@@ -63,6 +63,7 @@ pages = [
         "Contents" => "library/outline.md",
         "Particles" => "library/particles.md",
         "Phase Space Definition" => "library/phasespacedef.md",
+        "Phase Space Layout" => "library/phase_space_layout.md",
         "Phase Space Points" => "library/phasespacepoint.md",
         "Vector Types" => "library/vectors.md",
         "Lorentz Boosts" => "library/lorentzboosts.md",

docs/src/library/phase_space_layout.md

@@ -0,0 +1,29 @@
+# Phase Space Layout
+
+```@meta
+CurrentModule = QEDcore
+```
+
+## Phase Space Layout Interface
+
+```@docs
+AbstractPhaseSpaceLayout
+AbstractInPhaseSpaceLayout
+AbstractOutPhaseSpaceLayout
+_build_momenta
+phase_space_dimension
+in_phase_space_layout
+```
+
+## Building Momenta
+
+```@docs
+build_momenta
+```
+
+## Coordinate Maps
+
+```@docs
+CoordinateMap
+CoordinateMapCached
+```

src/QEDcore.jl

@@ -24,6 +24,10 @@ export Electron, Positron, Photon
 # particle base states
 export base_state
 
+# phase space layout
+export CoordinateMap
+export CoordinateMapCached
+
 # phase space
 export SphericalCoordinateSystem
 export CenterOfMomentumFrame, ElectronRestFrame
@@ -38,9 +42,14 @@ using SimpleTraits
 
 @reexport using QEDbase
 
+include("patch_QEDbase.jl")
+
 include("algebraic_objects/dirac_tensors/types.jl")
 include("algebraic_objects/dirac_tensors/multiplication.jl")
 
+include("coordinate_map/map.jl")
+include("coordinate_map/cached.jl")
+
 include("phase_spaces/types.jl")
 include("phase_spaces/access.jl")
 include("phase_spaces/create.jl")

src/coordinate_map/cached.jl

@@ -0,0 +1,112 @@
+"""
+    CoordinateMapCached{P,M,PSL,TM}(proc::P, model::M, psl::PSL, in_moms::TM)
+
+A `CoordinateMapCached` represents a precomputed transformation for phase space coordinates,
+where the momenta of the incoming particles are cached. This can improve performance in cases
+where the incoming momenta remain constant and only the outgoing momenta need to be computed
+based on new coordinates.
+
+The cached map encapsulates the scattering process (`proc`), the physics model (`model`),
+the phase space layout (`psl`), and the precomputed incoming momenta (`in_moms`).
+
+## Fields
+- `proc`: The scattering process definition, which defines the incoming and outgoing particles.
+- `model`: The physics model, which governs the interaction type and momentum distributions.
+- `psl`: The phase space layout, either incoming or outgoing, that maps coordinates to momenta.
+- `in_moms`: A collection of precomputed four-momenta for the incoming particles, usually a `Tuple`.
+
+## Usage
+
+### Cached Incoming Coordinates
+When a `CoordinateMapCached` build with `psl::AbstractInPhaseSpaceLayout` is called without
+any arguments, it returns the precomputed incoming momenta (`in_moms`) directly. This provides
+an efficient way to access the incoming particle momenta that have already been calculated and
+stored in the map.
+
+```julia
+# Create a coordinate map for in phase space
+in_coord_map = CoordinateMapCached(proc, model, in_psl, in_moms)
+
+# call to return the in-momenta
+in_coord_map()
+```
+
+- **Returns**: The cached tuple of four-momenta for the incoming particles.
+
+### Outgoing Coordinates
+The `CoordinateMapCached` can be called with outgoing phase space coordinates (`out_coords`).
+The cached incoming momenta (`in_moms`) are used to compute the total momentum (`Ptot`), which
+is then passed along with the outgoing coordinates to compute the momenta of the outgoing
+particles.
+
+```julia
+# Create a coordinate map for out phase space
+out_coord_map = CoordinateMapCached(proc, model, out_psl, in_moms)
+
+# call on out coordinates to return the out-momenta
+out_coord_map(out_coords)
+```
+
+- **Arguments**:
+    - `out_coords`: A tuple of phase space coordinates for the outgoing particles.
+
+- **Returns**:
+    - A tuple of four-momenta for the outgoing particles, consistent with the total momentum
+        derived from the cached incoming momenta.
+
+## Notes
+- **Caching**: The `CoordinateMapCached` is useful when the incoming momenta are fixed or do
+    not change frequently, as it avoids recomputation by storing the incoming momenta in the
+    cache.
+- **Efficiency**: This caching mechanism can significantly enhance performance when repeatedly
+    evaluating the scattering process with fixed incoming particles but varying outgoing
+    configurations.
+- The type is designed to handle both incoming and outgoing momenta, ensuring proper energy
+    and momentum conservation for the process.
+
+"""
+struct CoordinateMapCached{
+    P<:AbstractProcessDefinition,M<:AbstractModelDefinition,PSL<:AbstractPhaseSpaceLayout,TM
+} <: AbstractCoordinateMap
+    proc::P
+    model::M
+    psl::PSL
+    in_moms::TM
+end
+
+function CoordinateMapCached(
+    proc::AbstractProcessDefinition,
+    model::AbstractModelDefinition,
+    psl::AbstractInPhaseSpaceLayout,
+    in_coords::NTuple{N,T},
+) where {N,T<:Real}
+    in_moms = build_momenta(proc, model, psl, in_coords)
+    return CoordinateMapCached(proc, model, psl, in_moms)
+end
+
+function CoordinateMapCached(
+    proc::AbstractProcessDefinition,
+    model::AbstractModelDefinition,
+    psl::AbstractOutPhaseSpaceLayout,
+    in_coords::NTuple{N,T},
+) where {N,T<:Real}
+    in_moms = build_momenta(proc, model, in_phase_space_layout(psl), in_coords)
+    return CoordinateMapCached(proc, model, psl, in_moms)
+end
+
+# make the transform callable: for in_psl maps return the cached
+@inline function (
+    coord_map::CoordinateMapCached{P,M,PSL}
+)() where {P,M,PSL<:AbstractInPhaseSpaceLayout}
+    return getfield(coord_map, :in_moms)
+end
+
+# make the transform callable: for out_psl maps
+@inline function (coord_map::CoordinateMapCached{P,M,PSL})(
+    out_coords::Tuple
+) where {P,M,PSL<:AbstractOutPhaseSpaceLayout}
+    in_moms = coord_map.in_moms
+    return build_momenta(
+        coord_map.proc, coord_map.model, in_moms, coord_map.psl, out_coords
+    )
+end

src/coordinate_map/map.jl

@@ -0,0 +1,91 @@
+abstract type AbstractCoordinateMap end
+
+"""
+    CoordinateMap{P,M,PSL}(proc::P, model::M, psl::PSL)
+
+A `CoordinateMap` represents a transformation that maps phase space coordinates to particle
+momenta for a specific scattering process. This type encapsulates the scattering process
+definition (`proc`), the physics model (`model`), and the phase space layout (`psl`).
+The `CoordinateMap` is callable and supports the conversion of both incoming and outgoing
+phase space coordinates into four-momenta.
+
+## Fields
+- `proc`: The scattering process definition, subtype of `AbstractProcessDefinition`.
+- `model`: The physics model, subtype of `AbstractModelDefinition`.
+- `psl`: The phase space layout, either incoming or outgoing, that maps coordinates to momenta.
+
+## Usage
+
+### Incoming Coordinates
+The `CoordinateMap` can be called with a tuple of incoming phase space coordinates
+(`in_coords`), which are used to compute the corresponding incoming particle four-momenta.
+This is done by calling the phase space construction function [`build_momenta`](@ref) using
+the `proc`, `model`, and `psl` provided in the `CoordinateMap`.
+
+```julia
+# Create a coordinate map for in phase space
+in_coord_map = CoordinateMap(proc, model, in_psl)
+
+# call on in-coordinates to build the momenta
+in_coord_map(in_coords)
+```
+
+- **Arguments**:
+    - `in_coords`: A tuple of phase space coordinates for the incoming particles.
+
+- **Returns**: A tuple of four-momenta corresponding to the incoming particles.
+
+### Incoming and Outgoing Coordinates
+The `CoordinateMap` can also be called with both incoming (`in_coords`) and outgoing
+(`out_coords`) phase space coordinates. The incoming momenta are computed first, and then
+the total momentum (`Ptot`) is calculated by summing these momenta. This total momentum is
+then used to compute the outgoing particle four-momenta using the provided outgoing phase
+space layout.
+
+```julia
+# Create a coordinate map for out phase space
+coord_map = CoordinateMap(proc, model, out_psl)
+
+# call on in- and out-coordinates to build the momenta
+coord_map(in_coords, out_coords)
+```
+
+- **Arguments**:
+    - `in_coords`: A tuple of phase space coordinates for the incoming particles.
+    - `out_coords`: A tuple of phase space coordinates for the outgoing particles.
+
+- **Returns**:
+    - A tuple where the first element is the incoming particle momenta and the second element
+        is the outgoing particle momenta, both represented as tuples of four-momenta.
+
+## Notes
+- The `CoordinateMap` provides a flexible mechanism for transforming phase space coordinates
+    into physically meaningful four-momenta, ensuring consistency with the scattering process,
+    physics model, and phase space layout.
+- The type is designed to handle both incoming and outgoing momenta, ensuring proper energy
+    and momentum conservation for the process.
+
+"""
+struct CoordinateMap{P,M,PSL<:AbstractPhaseSpaceLayout} <: AbstractCoordinateMap
+    proc::P
+    model::M
+    psl::PSL
+end
+
+# make the transform callable
+@inline function (coord_map::CoordinateMap{P,M,PSL})(
+    in_coords::Tuple
+) where {P,M,PSL<:AbstractInPhaseSpaceLayout}
+    return build_momenta(coord_map.proc, coord_map.model, coord_map.psl, in_coords)
+end
+
+# make the transform callable
+@inline function (coord_map::CoordinateMap{P,M,PSL})(
+    in_coords::Tuple, out_coords::Tuple
+) where {P,M,PSL<:AbstractOutPhaseSpaceLayout}
+    in_moms = build_momenta(
+        coord_map.proc, coord_map.model, in_phase_space_layout(coord_map.psl), in_coords
+    )
+    return in_moms,
+    build_momenta(coord_map.proc, coord_map.model, in_moms, coord_map.psl, out_coords)
+end