Line data    Source code

       1           0 : // Distributed under the MIT License.
       2             : // See LICENSE.txt for details.
       3             : 
       4             : #pragma once
       5             : 
       6             : #include <cstddef>
       7             : #include <memory>
       8             : #include <pup.h>
       9             : #include <tuple>
      10             : #include <type_traits>
      11             : #include <unordered_map>
      12             : #include <unordered_set>
      13             : #include <utility>
      14             : #include <vector>
      15             : 
      16             : #include "DataStructures/Tensor/Tensor.hpp"
      17             : #include "Domain/CoordinateMaps/CoordinateMap.hpp"
      18             : #include "Domain/FunctionsOfTime/FunctionOfTime.hpp"
      19             : #include "Utilities/Serialization/CharmPupable.hpp"
      20             : #include "Utilities/TMPL.hpp"
      21             : #include "Utilities/Autodiff/Autodiff.hpp"
      22             : #include "Utilities/Simd/Simd.hpp"
      23             : 
      24             : namespace domain::CoordinateMaps {
      25             : 
      26             : /*!
      27             :  * \brief A composition of coordinate maps at runtime
      28             :  *
      29             :  * Composes a sequence of `domain::CoordinateMapBase` that step through the
      30             :  * `Frames`. The result is another `domain::CoordinateMapBase`. This is
      31             :  * different to `domain::CoordinateMap`, which does the composition at compile
      32             :  * time. The use cases are different:
      33             :  *
      34             :  * - Use `domain::CoordinateMap` to compose maps at compile time to go from one
      35             :  *   (named) frame to another using any number of coordinate transformation. The
      36             :  *   coordinate transformations are concatenated to effectively form a single
      37             :  *   map, and intermediate frames have no meaning. This has the performance
      38             :  *   benefit that looking up pointers and calling into virtual member functions
      39             :  *   of intermediate maps is avoided, and it has the semantic benefit that
      40             :  *   intermediate frames without meaning are not named or even accessible.
      41             :  *   **Example:** A static BlockLogical -> Grid map that deforms the logical
      42             :  *   cube to a wedge, applies a radial redistribution of grid points, and
      43             :  *   translates + rotates the wedge.
      44             :  * - Use `domain::CoordinateMaps::Composition` (this class) to compose maps at
      45             :  *   runtime to step through a sequence of (named) frames.
      46             :  *   **Example:** A time-dependent ElementLogical -> BlockLogical -> Grid ->
      47             :  *   Inertial map that applies an affine transformation between the
      48             :  *   ElementLogical and BlockLogical frames (see
      49             :  *   domain::element_to_block_logical_map), then deforms the logical cube to a
      50             :  *   wedge using the map described above (Grid frame), and then rotates the grid
      51             :  *   with a time-dependent rotation map (Inertial frame).
      52             :  *
      53             :  * \warning Think about performance implications before using this Composition
      54             :  * class. In an evolution it's usually advantageous to keep at least some of the
      55             :  * maps separate, e.g. to avoid reevaluating the static maps in every time step.
      56             :  * You can also access individual components of the composition in this class.
      57             :  *
      58             :  * \tparam Frames The list of frames in the composition, as a tmpl::list<>.
      59             :  * The first entry in the list is the source frame, and the last entry is the
      60             :  * target frame of the composition. Maps in the composition step through the
      61             :  * `Frames`. For example, if `Frames = tmpl::list<Frame::ElementLogical,
      62             :  * Frame::BlockLogical, Frame::Inertial>`, then the composition has two maps:
      63             :  * ElementLogical -> BlockLogical and BlockLogical -> Inertial.
      64             :  *
      65             :  * ## Note on inverse Hessian computation
      66             :  * Below we work out the algebra for computing the inverse Hessian
      67             :  * used in this struct (if SPECTRE_AUTODIFF=ON). Suppose we have
      68             :  * a composition of maps \f$ \xi^i \longrightarrow \cdots \longrightarrow x^i
      69             :  * \longrightarrow y^i \f$, where \f$ \xi^i \f$ are the source coordinates, \f$
      70             :  * x^i \f$ are some intermediate coordinates, and \f$ y^i \f$ are the final
      71             :  * target coordinates,  We compute inverse Hessian by first
      72             :  * propagating autodiff dual types through the composed `call_impl` function to
      73             :  * automatically get the Hessian \f$ \frac{\partial^2 y^i}{\partial\xi^j
      74             :  * \partial\xi^k} \f$ Then the inverse Hessian is computed by the identity
      75             :  * \f[ \frac{\partial^2\xi^i}{\partial y^m \partial y^n} =
      76             :  * -\frac{\partial\xi^i}{\partial y^j}\frac{\partial\xi^k}{\partial y^m}
      77             :  * \frac{\partial\xi^l}{\partial y^n}\frac{\partial^2 y^j}{\partial\xi^k
      78             :  * \partial\xi^l}, \f] where the inverse Jacobian is passed in as a function
      79             :  * argument.
      80             :  *
      81             :  * See Test_Composition.cpp for a different implementation. The current
      82             :  * implementation is chosen in production as it is faster when we have
      83             :  * the inverse Jacobian already, and in most cases we do.
      84             :  *
      85             :  * ## Note on auto differentiation
      86             :  * We use forward mode autodiff here as it is simpler to implement and
      87             :  * has better optimization than the reverse mode. Reverse mode in
      88             :  * the [Autodiff](https://github.com/autodiff/autodiff) library
      89             :  * has higher cost per propagation and does not support taping.
      90             :  * Also see this
      91             :  * [github issue](https://github.com/autodiff/autodiff/issues/332).
      92             :  */
      93             : template <typename Frames, size_t Dim,
      94             :           typename Is = std::make_index_sequence<tmpl::size<Frames>::value - 1>>
      95           1 : struct Composition;
      96             : 
      97             : template <typename Frames, size_t Dim, size_t... Is>
      98             : struct Composition<Frames, Dim, std::index_sequence<Is...>>
      99             :     : public CoordinateMapBase<tmpl::front<Frames>, tmpl::back<Frames>, Dim> {
     100             :   using frames = Frames;
     101             :   using SourceFrame = tmpl::front<Frames>;
     102             :   using TargetFrame = tmpl::back<Frames>;
     103             :   static constexpr size_t num_frames = tmpl::size<Frames>::value;
     104             :   using Base = CoordinateMapBase<SourceFrame, TargetFrame, Dim>;
     105             :   using FuncOfTimeMap = std::unordered_map<
     106             :       std::string, std::unique_ptr<domain::FunctionsOfTime::FunctionOfTime>>;
     107             :   using Base::operator();
     108             : 
     109             :   Composition() = default;
     110             :   Composition(const Composition& rhs) { *this = rhs; }
     111             :   Composition& operator=(const Composition& rhs);
     112             :   Composition(Composition&& /*rhs*/) = default;
     113             :   Composition& operator=(Composition&& /*rhs*/) = default;
     114             :   ~Composition() override = default;
     115             : 
     116             :   std::unique_ptr<Base> get_clone() const override {
     117             :     return std::make_unique<Composition>(*this);
     118             :   }
     119             : 
     120             :   /// \cond
     121             :   explicit Composition(CkMigrateMessage* /*m*/) {}
     122             :   using PUP::able::register_constructor;
     123             :   WRAPPED_PUPable_decl_template(Composition);  // NOLINT
     124             :   /// \endcond
     125             : 
     126             :   Composition(std::unique_ptr<CoordinateMapBase<
     127             :                   tmpl::at<frames, tmpl::size_t<Is>>,
     128             :                   tmpl::at<frames, tmpl::size_t<Is + 1>>, Dim>>... maps);
     129             : 
     130             :   const std::tuple<std::unique_ptr<
     131             :       CoordinateMapBase<tmpl::at<frames, tmpl::size_t<Is>>,
     132             :                         tmpl::at<frames, tmpl::size_t<Is + 1>>, Dim>>...>&
     133             :   maps() const {
     134             :     return maps_;
     135             :   }
     136             : 
     137             :   template <typename LocalSourceFrame, typename LocalTargetFrame>
     138             :   const CoordinateMapBase<LocalSourceFrame, LocalTargetFrame, Dim>&
     139             :   get_component() const {
     140             :     return *get<std::unique_ptr<
     141             :         CoordinateMapBase<LocalSourceFrame, LocalTargetFrame, Dim>>>(maps_);
     142             :   }
     143             : 
     144             :   bool is_identity() const override;
     145             : 
     146             :   bool inv_jacobian_is_time_dependent() const override;
     147             : 
     148             :   bool jacobian_is_time_dependent() const override;
     149             : 
     150             :   const std::unordered_set<std::string>& function_of_time_names()
     151             :       const override {
     152             :     return function_of_time_names_;
     153             :   }
     154             : 
     155             :   bool supports_hessian() const override;
     156             : 
     157             :   tnsr::I<double, Dim, TargetFrame> operator()(
     158             :       tnsr::I<double, Dim, SourceFrame> source_point,
     159             :       double time = std::numeric_limits<double>::signaling_NaN(),
     160             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     161             : 
     162             :   tnsr::I<DataVector, Dim, TargetFrame> operator()(
     163             :       tnsr::I<DataVector, Dim, SourceFrame> source_point,
     164             :       double time = std::numeric_limits<double>::signaling_NaN(),
     165             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     166             : 
     167             :   std::optional<tnsr::I<double, Dim, SourceFrame>> inverse(
     168             :       tnsr::I<double, Dim, TargetFrame> target_point,
     169             :       double time = std::numeric_limits<double>::signaling_NaN(),
     170             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     171             : 
     172             :   InverseJacobian<double, Dim, SourceFrame, TargetFrame> inv_jacobian(
     173             :       tnsr::I<double, Dim, SourceFrame> source_point,
     174             :       double time = std::numeric_limits<double>::signaling_NaN(),
     175             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     176             : 
     177             :   InverseJacobian<DataVector, Dim, SourceFrame, TargetFrame> inv_jacobian(
     178             :       tnsr::I<DataVector, Dim, SourceFrame> source_point,
     179             :       double time = std::numeric_limits<double>::signaling_NaN(),
     180             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     181             : 
     182             : #ifdef SPECTRE_AUTODIFF
     183             :   InverseJacobian<autodiff::HigherOrderDual<2, double>, Dim, SourceFrame,
     184             :                   TargetFrame>
     185             :   inv_jacobian(tnsr::I<autodiff::HigherOrderDual<2, double>, Dim, SourceFrame>
     186             :                    source_point,
     187             :                double time = std::numeric_limits<double>::signaling_NaN(),
     188             :                const FuncOfTimeMap& functions_of_time = {}) const override;
     189             : 
     190             :   // Compute inverse hessian by calling operator()
     191             :   InverseHessian<double, Dim, SourceFrame, TargetFrame> inv_hessian(
     192             :       tnsr::I<double, Dim, SourceFrame> source_point,
     193             :       const InverseJacobian<double, Dim, SourceFrame, TargetFrame>& inverse_jac,
     194             :       double time = std::numeric_limits<double>::signaling_NaN(),
     195             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     196             : 
     197             :   // Compute inverse hessian by calling operator()
     198             :   InverseHessian<DataVector, Dim, SourceFrame, TargetFrame> inv_hessian(
     199             :       tnsr::I<DataVector, Dim, SourceFrame> source_point,
     200             :       const InverseJacobian<DataVector, Dim, SourceFrame, TargetFrame>&
     201             :           inverse_jac,
     202             :       double time = std::numeric_limits<double>::signaling_NaN(),
     203             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     204             : #endif  // SPECTRE_AUTODIFF
     205             : 
     206             :   Jacobian<double, Dim, SourceFrame, TargetFrame> jacobian(
     207             :       tnsr::I<double, Dim, SourceFrame> source_point,
     208             :       double time = std::numeric_limits<double>::signaling_NaN(),
     209             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     210             : 
     211             :   Jacobian<DataVector, Dim, SourceFrame, TargetFrame> jacobian(
     212             :       tnsr::I<DataVector, Dim, SourceFrame> source_point,
     213             :       double time = std::numeric_limits<double>::signaling_NaN(),
     214             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     215             : 
     216             :   [[noreturn]] std::tuple<
     217             :       tnsr::I<double, Dim, TargetFrame>,
     218             :       InverseJacobian<double, Dim, SourceFrame, TargetFrame>,
     219             :       Jacobian<double, Dim, SourceFrame, TargetFrame>,
     220             :       tnsr::I<double, Dim, TargetFrame>>
     221             :   coords_frame_velocity_jacobians(
     222             :       tnsr::I<double, Dim, SourceFrame> source_point,
     223             :       double time = std::numeric_limits<double>::signaling_NaN(),
     224             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     225             : 
     226             :   [[noreturn]] std::tuple<
     227             :       tnsr::I<DataVector, Dim, TargetFrame>,
     228             :       InverseJacobian<DataVector, Dim, SourceFrame, TargetFrame>,
     229             :       Jacobian<DataVector, Dim, SourceFrame, TargetFrame>,
     230             :       tnsr::I<DataVector, Dim, TargetFrame>>
     231             :   coords_frame_velocity_jacobians(
     232             :       tnsr::I<DataVector, Dim, SourceFrame> source_point,
     233             :       double time = std::numeric_limits<double>::signaling_NaN(),
     234             :       const FuncOfTimeMap& functions_of_time = {}) const override;
     235             : 
     236             :   [[noreturn]] std::unique_ptr<CoordinateMapBase<SourceFrame, Frame::Grid, Dim>>
     237             :   get_to_grid_frame() const override;
     238             : 
     239             :   // NOLINTNEXTLINE(google-runtime-references)
     240             :   void pup(PUP::er& p) override;
     241             : 
     242             :  private:
     243             :   template <typename DataType>
     244             :   tnsr::I<DataType, Dim, TargetFrame> call_impl(
     245             :       tnsr::I<DataType, Dim, SourceFrame> source_point,
     246             :       double time = std::numeric_limits<double>::signaling_NaN(),
     247             :       const FuncOfTimeMap& functions_of_time = {}) const;
     248             : 
     249             :   template <typename DataType>
     250             :   std::optional<tnsr::I<DataType, Dim, SourceFrame>> inverse_impl(
     251             :       tnsr::I<DataType, Dim, TargetFrame> target_point,
     252             :       double time = std::numeric_limits<double>::signaling_NaN(),
     253             :       const FuncOfTimeMap& functions_of_time = {}) const;
     254             : 
     255             :   template <typename DataType>
     256             :   InverseJacobian<DataType, Dim, SourceFrame, TargetFrame> inv_jacobian_impl(
     257             :       tnsr::I<DataType, Dim, SourceFrame> source_point,
     258             :       double time = std::numeric_limits<double>::signaling_NaN(),
     259             :       const FuncOfTimeMap& functions_of_time = {}) const;
     260             : 
     261             :   template <typename DataType>
     262             :   InverseHessian<DataType, Dim, SourceFrame, TargetFrame> inv_hessian_impl(
     263             :       tnsr::I<DataType, Dim, SourceFrame> source_point,
     264             :       const ::InverseJacobian<DataType, Dim, SourceFrame, TargetFrame>&
     265             :           inverse_jac,
     266             :       double time = std::numeric_limits<double>::signaling_NaN(),
     267             :       const FuncOfTimeMap& functions_of_time = {}) const;
     268             : 
     269             :   template <typename DataType>
     270             :   Jacobian<DataType, Dim, SourceFrame, TargetFrame> jacobian_impl(
     271             :       tnsr::I<DataType, Dim, SourceFrame> source_point,
     272             :       double time = std::numeric_limits<double>::signaling_NaN(),
     273             :       const FuncOfTimeMap& functions_of_time = {}) const;
     274             : 
     275             :   bool is_equal_to(const CoordinateMapBase<SourceFrame, TargetFrame, Dim>&
     276             :                        other) const override;
     277             : 
     278             :   std::tuple<std::unique_ptr<
     279             :       CoordinateMapBase<tmpl::at<frames, tmpl::size_t<Is>>,
     280             :                         tmpl::at<frames, tmpl::size_t<Is + 1>>, Dim>>...>
     281             :       maps_;
     282             :   std::unordered_set<std::string> function_of_time_names_;
     283             : };
     284             : 
     285             : // Template deduction guide
     286             : template <typename FirstMap, typename... Maps>
     287           0 : Composition(std::unique_ptr<FirstMap>, std::unique_ptr<Maps>... maps)
     288             :     -> Composition<tmpl::list<typename FirstMap::source_frame,
     289             :                               typename FirstMap::target_frame,
     290             :                               typename Maps::target_frame...>,
     291             :                    FirstMap::dim>;
     292             : 
     293             : /// \cond
     294             : template <typename Frames, size_t Dim, size_t... Is>
     295             : PUP::able::PUP_ID
     296             :     Composition<Frames, Dim, std::index_sequence<Is...>>::my_PUP_ID = 0;
     297             : /// \endcond
     298             : 
     299             : }  // namespace domain::CoordinateMaps