Add optical track initialization action (#1438)

* Add optical track initialization action

* Add optical sim data/view

* Add optical particle data/view

* Add initialization action to optical loop and fix offset error in generators

* Clean up

src/celeritas/CMakeLists.txt

@@ -283,6 +283,7 @@ celeritas_polysource(global/detail/TrackSlotUtils)
 celeritas_polysource(neutron/model/ChipsNeutronElasticModel)
 celeritas_polysource(neutron/model/NeutronInelasticModel)
 celeritas_polysource(optical/action/BoundaryAction)
+celeritas_polysource(optical/action/InitializeTracksAction)
 celeritas_polysource(optical/detail/CerenkovGeneratorAction)
 celeritas_polysource(optical/detail/CerenkovOffloadAction)
 celeritas_polysource(optical/detail/OpticalGenAlgorithms)

src/celeritas/optical/CoreParams.cc

@@ -20,6 +20,7 @@
 #include "MaterialParams.hh"
 #include "TrackInitParams.hh"
 #include "action/BoundaryAction.hh"
+#include "action/InitializeTracksAction.hh"
 
 namespace celeritas
 {
@@ -44,7 +45,6 @@ build_params_refs(CoreParams::Input const& p, CoreScalars const& scalars)
     ref.material = get_ref<M>(*p.material);
     // TODO: ref.physics = get_ref<M>(*p.physics);
     ref.rng = get_ref<M>(*p.rng);
-    ref.sim = get_ref<M>(*p.sim);
     ref.init = get_ref<M>(*p.init);
 
     CELER_ENSURE(ref);
@@ -63,9 +63,7 @@ CoreScalars build_actions(ActionRegistry* reg)
 
     //// START ACTIONS ////
 
-#if 0
     reg->insert(make_shared<InitializeTracksAction>(reg->next_id()));
-#endif
 
     //// PRE-STEP ACTIONS ////
 
@@ -101,7 +99,6 @@ CoreParams::CoreParams(Input&& input) : input_(std::move(input))
     CP_VALIDATE_INPUT(material);
     // TODO: CP_VALIDATE_INPUT(physics);
     CP_VALIDATE_INPUT(rng);
-    CP_VALIDATE_INPUT(sim);
     CP_VALIDATE_INPUT(init);
     CP_VALIDATE_INPUT(action_reg);
     CP_VALIDATE_INPUT(max_streams);

src/celeritas/optical/CoreParams.hh

@@ -20,7 +20,6 @@ namespace celeritas
 {
 //---------------------------------------------------------------------------//
 class ActionRegistry;
-class SimParams;
 
 namespace optical
 {
@@ -41,7 +40,6 @@ class CoreParams final : public ParamsDataInterface<CoreParamsData>
     using SPConstGeo = std::shared_ptr<GeoParams const>;
     using SPConstMaterial = std::shared_ptr<MaterialParams const>;
     using SPConstRng = std::shared_ptr<RngParams const>;
-    using SPConstSim = std::shared_ptr<SimParams const>;
     using SPConstTrackInit = std::shared_ptr<TrackInitParams const>;
     using SPActionRegistry = std::shared_ptr<ActionRegistry>;
 
@@ -57,7 +55,6 @@ class CoreParams final : public ParamsDataInterface<CoreParamsData>
         SPConstMaterial material;
         // TODO: physics
         SPConstRng rng;
-        SPConstSim sim;
         SPConstTrackInit init;
 
         SPActionRegistry action_reg;
@@ -68,7 +65,7 @@ class CoreParams final : public ParamsDataInterface<CoreParamsData>
         //! True if all params are assigned and valid
         explicit operator bool() const
         {
-            return geometry && material && rng && sim && init && action_reg
+            return geometry && material && rng && init && action_reg
                    && max_streams;
         }
     };

src/celeritas/optical/CoreTrackData.cc

@@ -35,9 +35,10 @@ void resize(CoreStateData<Ownership::value, M>* state,
     // Geant4 state is stream-local
     resize(&state->geometry, params.geometry, stream_id, size);
 #endif
+    resize(&state->particle, size);
     resize(&state->physics, params.physics, size);
     resize(&state->rng, params.rng, stream_id, size);
-    resize(&state->sim, params.sim, size);
+    resize(&state->sim, size);
     resize(&state->init, params.init, stream_id, size);
     state->stream_id = stream_id;
 

src/celeritas/optical/CoreTrackData.hh

@@ -12,10 +12,11 @@
 #include "celeritas/Types.hh"
 #include "celeritas/geo/GeoData.hh"
 #include "celeritas/random/RngData.hh"
-#include "celeritas/track/SimData.hh"
 
 #include "CoreTrackDataFwd.hh"
 #include "MaterialData.hh"
+#include "ParticleData.hh"
+#include "SimData.hh"
 #include "TrackInitData.hh"
 #include "Types.hh"
 
@@ -83,15 +84,14 @@ struct CoreParamsData
     MaterialParamsData<W, M> material;
     PhysicsParamsData<W, M> physics;
     RngParamsData<W, M> rng;
-    SimParamsData<W, M> sim;
     TrackInitParamsData<W, M> init;
 
     CoreScalars scalars;
 
     //! True if all params are assigned
     explicit CELER_FUNCTION operator bool() const
     {
-        return geometry && material && physics && rng && sim && init && scalars;
+        return geometry && material && physics && rng && init && scalars;
     }
 
     //! Assign from another set of data
@@ -103,7 +103,6 @@ struct CoreParamsData
         material = other.material;
         physics = other.physics;
         rng = other.rng;
-        sim = other.sim;
         init = other.init;
         scalars = other.scalars;
         return *this;
@@ -122,9 +121,10 @@ struct CoreStateData
 
     GeoStateData<W, M> geometry;
     // TODO: should we cache the material ID?
+    ParticleStateData<W, M> particle;
     PhysicsStateData<W, M> physics;
     RngStateData<W, M> rng;
-    SimStateData<W, M> sim;  // TODO: has a few things we don't need
+    SimStateData<W, M> sim;
     TrackInitStateData<W, M> init;
 
     //! Unique identifier for "thread-local" data.
@@ -136,7 +136,8 @@ struct CoreStateData
     //! Whether the data are assigned
     explicit CELER_FUNCTION operator bool() const
     {
-        return geometry && physics && rng && sim && init && stream_id;
+        return geometry && particle && physics && rng && sim && init
+               && stream_id;
     }
 
     //! Assign from another set of data
@@ -145,6 +146,7 @@ struct CoreStateData
     {
         CELER_EXPECT(other);
         geometry = other.geometry;
+        particle = other.particle;
         physics = other.physics;
         rng = other.rng;
         sim = other.sim;

src/celeritas/optical/CoreTrackView.hh

@@ -9,11 +9,16 @@
 
 #include "celeritas/geo/GeoTrackView.hh"
 #include "celeritas/random/RngEngine.hh"
-#include "celeritas/track/SimTrackView.hh"
 
 #include "CoreTrackData.hh"
 #include "MaterialView.hh"
-#include "TrackView.hh"
+#include "ParticleTrackView.hh"
+#include "SimTrackView.hh"
+#include "TrackInitializer.hh"
+
+#if !CELER_DEVICE_COMPILE
+#    include "corecel/io/Logger.hh"
+#endif
 
 namespace celeritas
 {
@@ -38,6 +43,9 @@ class CoreTrackView
                                         StateRef const& states,
                                         TrackSlotId slot);
 
+    // Initialize the track states
+    inline CELER_FUNCTION CoreTrackView& operator=(TrackInitializer const&);
+
     // Return a geometry view
     inline CELER_FUNCTION GeoTrackView geometry() const;
 
@@ -50,8 +58,8 @@ class CoreTrackView
     // Return a simulation management view
     inline CELER_FUNCTION SimTrackView sim() const;
 
-    // Return a physics view
-    inline CELER_FUNCTION TrackView physics() const;
+    // Return a particle view
+    inline CELER_FUNCTION ParticleTrackView particle() const;
 
     // Return an RNG engine
     inline CELER_FUNCTION RngEngine rng() const;
@@ -88,6 +96,42 @@ CoreTrackView::CoreTrackView(ParamsRef const& params,
     CELER_EXPECT(track_slot_id_ < states_.size());
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Initialize the track states.
+ */
+CELER_FUNCTION CoreTrackView&
+CoreTrackView::operator=(TrackInitializer const& init)
+{
+    // Initialiize the sim state
+    this->sim() = SimTrackView::Initializer{init.time};
+
+    // Initialize the geometry state
+    auto geo = this->geometry();
+    geo = GeoTrackInitializer{init.position, init.direction};
+    if (CELER_UNLIKELY(geo.failed() || geo.is_outside()))
+    {
+#if !CELER_DEVICE_COMPILE
+        if (geo.is_outside())
+        {
+            // Print an error message if initialization was "successful" but
+            // track is outside
+            CELER_LOG_LOCAL(error) << "Track started outside the geometry";
+        }
+#endif
+        this->apply_errored();
+        return *this;
+    }
+
+    // Initialize the particle state
+    this->particle()
+        = ParticleTrackView::Initializer{init.energy, init.polarization};
+
+    //! \todo Add physics view and clear physics state
+
+    return *this;
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Return a geometry view.
@@ -120,11 +164,11 @@ CoreTrackView::material(GeoTrackView const& geo) const -> MaterialView
 
 //---------------------------------------------------------------------------//
 /*!
- * Return a physics view.
+ * Return a particle view.
  */
-CELER_FUNCTION auto CoreTrackView::physics() const -> TrackView
+CELER_FUNCTION auto CoreTrackView::particle() const -> ParticleTrackView
 {
-    CELER_NOT_IMPLEMENTED("physics track view");
+    return ParticleTrackView{states_.particle, this->track_slot_id()};
 }
 
 //---------------------------------------------------------------------------//
@@ -142,7 +186,7 @@ CELER_FUNCTION auto CoreTrackView::rng() const -> RngEngine
  */
 CELER_FUNCTION SimTrackView CoreTrackView::sim() const
 {
-    return SimTrackView{params_.sim, states_.sim, this->track_slot_id()};
+    return SimTrackView{states_.sim, this->track_slot_id()};
 }
 
 //---------------------------------------------------------------------------//

src/celeritas/optical/OpticalCollector.cc

@@ -10,7 +10,6 @@
 #include "corecel/data/AuxParamsRegistry.hh"
 #include "corecel/sys/ActionRegistry.hh"
 #include "celeritas/global/CoreParams.hh"
-#include "celeritas/track/SimParams.hh"
 #include "celeritas/track/TrackInitParams.hh"
 
 #include "CerenkovParams.hh"

src/celeritas/optical/ParticleData.hh

@@ -0,0 +1,76 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/optical/ParticleData.hh
+//---------------------------------------------------------------------------//
+#pragma once
+
+#include "corecel/Macros.hh"
+#include "corecel/Types.hh"
+#include "corecel/cont/Array.hh"
+#include "corecel/data/Collection.hh"
+#include "corecel/data/CollectionBuilder.hh"
+#include "celeritas/Types.hh"
+
+namespace celeritas
+{
+namespace optical
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Storage for dynamic particle data.
+ */
+template<Ownership W, MemSpace M>
+struct ParticleStateData
+{
+    //// TYPES ////
+
+    using Real3 = Array<real_type, 3>;
+    template<class T>
+    using Items = celeritas::StateCollection<T, W, M>;
+
+    //// DATA ////
+
+    Items<real_type> energy;  //! Kinetic energy [MeV]
+    Items<Real3> polarization;
+
+    //// METHODS ////
+
+    //! Check whether the interface is assigned
+    explicit CELER_FUNCTION operator bool() const
+    {
+        return !energy.empty() && !polarization.empty();
+    }
+
+    //! State size
+    CELER_FUNCTION size_type size() const { return energy.size(); }
+
+    //! Assign from another set of data
+    template<Ownership W2, MemSpace M2>
+    ParticleStateData& operator=(ParticleStateData<W2, M2>& other)
+    {
+        CELER_EXPECT(other);
+        energy = other.energy;
+        polarization = other.polarization;
+        return *this;
+    }
+};
+
+//---------------------------------------------------------------------------//
+/*!
+ * Resize particle states.
+ */
+template<MemSpace M>
+inline void resize(ParticleStateData<Ownership::value, M>* data, size_type size)
+{
+    CELER_EXPECT(size > 0);
+    resize(&data->energy, size);
+    resize(&data->polarization, size);
+    CELER_ENSURE(*data);
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace optical
+}  // namespace celeritas