Line data    Source code

       1             : //
       2             : // Initialize a field using a mathematical expression.
       3             : // Expressions are imported as strings and are compiled real-time using the
       4             : // `AMReX Parser <https://amrex-codes.github.io/amrex/docs_html/Basics.html#parser>`_.
       5             : //
       6             : // Works for single or multiple-component fields.
       7             : // Use the :code:`regionN` (N=0,1,2, etc. up to number of components) to pass expression.
       8             : // For example:
       9             : //
      10             : // .. code-block:: bash
      11             : // 
      12             : //    ic.region0 = "sin(x*y*z)"
      13             : //    ic.region1 = "3.0*(x > 0.5 and y > 0.5)"
      14             : //
      15             : // for a two-component field. It is up to you to make sure your expressions are parsed
      16             : // correctly; otherwise you will get undefined behavior.
      17             : //
      18             : // :bdg-primary-line:`Constants`
      19             : // You can add constants to your expressions using the :code:`constant` directive.
      20             : // For instance, in the following code
      21             : // 
      22             : // .. code-block:: bash
      23             : //
      24             : //    psi.ic.type=expression
      25             : //    psi.ic.expression.constant.eps = 0.05
      26             : //    psi.ic.expression.constant.R   = 0.25
      27             : //    psi.ic.expression.region0 = "0.5 + 0.5*tanh((x^2 + y^2 - R)/eps)"
      28             : //    
      29             : // the constants :code:`eps` and :code:`R` are defined by the user and then used
      30             : // in the subsequent expression.
      31             : // The variables can have any name made up of characters that is not reserved.
      32             : // However, if multiple ICs are used, they must be defined each time for each IC.
      33             : //
      34             : 
      35             : #ifndef IC_EXPRESSION_H_
      36             : #define IC_EXPRESSION_H_
      37             : #include "IC/IC.H"
      38             : #include "Util/Util.H"
      39             : #include "IO/ParmParse.H"
      40             : #include "AMReX_Parser.H"
      41             : 
      42             : namespace IC
      43             : {
      44             : class Expression : public IC
      45             : {
      46             : private:
      47             :     enum CoordSys { Cartesian, Polar, Spherical };
      48             :     std::vector<amrex::Parser> parser;
      49             :     std::vector<amrex::ParserExecutor<4>> f;
      50             :     Expression::CoordSys coord = Expression::CoordSys::Cartesian;
      51             : public:
      52             :     static constexpr const char* name = "expression";
      53             :     Expression(amrex::Vector<amrex::Geometry>& _geom) : IC(_geom) {}
      54           3 :     Expression(amrex::Vector<amrex::Geometry>& _geom, IO::ParmParse& pp, std::string name) : IC(_geom)
      55             :     {
      56           3 :         pp_queryclass(name, *this);
      57           3 :     }
      58          20 :     void Add(const int& lev, Set::Field<Set::Scalar>& a_field, Set::Scalar a_time = 0.0)
      59             :     {
      60          40 :         Util::Assert(INFO, TEST(a_field[lev]->nComp() == (int)f.size()));
      61         105 :         for (amrex::MFIter mfi(*a_field[lev], amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
      62             :         {
      63          85 :             amrex::Box bx;// = mfi.tilebox();
      64             :             //bx.grow(a_field[lev]->nGrow());
      65          85 :             amrex::IndexType type = a_field[lev]->ixType();
      66         170 :             if (type == amrex::IndexType::TheCellType())      bx = mfi.growntilebox();
      67           0 :             else if (type == amrex::IndexType::TheNodeType()) bx = mfi.grownnodaltilebox();
      68           0 :             else Util::Abort(INFO, "Unkonwn index type");
      69             : 
      70          85 :             amrex::Array4<Set::Scalar> const& field = a_field[lev]->array(mfi);
      71         170 :             for (unsigned int n = 0; n < f.size(); n++)
      72             :             {
      73          85 :                 amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k)
      74             :                 {
      75       29120 :                     Set::Vector x = Set::Position(i, j, k, geom[lev], type);
      76       29120 :                     if (coord == Expression::CoordSys::Cartesian)
      77             :                     {
      78             : #if AMREX_SPACEDIM == 1
      79             :                         field(i, j, k, n) = f[n](x(0), 0.0, 0.0, a_time);
      80             : #elif AMREX_SPACEDIM == 2
      81       87360 :                         field(i, j, k, n) = f[n](x(0), x(1), 0.0, a_time);
      82             : #elif AMREX_SPACEDIM == 3
      83           0 :                         field(i, j, k, n) = f[n](x(0), x(1), x(2), a_time);
      84             : #endif
      85             :                     }
      86             : #if AMREX_SPACEDIM>1
      87           0 :                     else if (coord == Expression::CoordSys::Polar)
      88             :                     {
      89           0 :                         field(i, j, k, n) = f[n](sqrt(x(0)* x(0) + x(1) * x(1)), std::atan2(x(1), x(0)), x(2), a_time);
      90             :                     }
      91             : #endif
      92       29120 :                 });
      93             :             }
      94             :         }
      95          20 :         a_field[lev]->FillBoundary();
      96          20 :     };
      97             : 
      98           0 :     void Add(const int& lev, Set::Field<Set::Vector>& a_field, Set::Scalar a_time = 0.0)
      99             :     {
     100           0 :         Util::Assert(INFO, TEST(a_field[lev]->nComp() == 1));
     101           0 :         Util::Assert(INFO, TEST(f.size() >= AMREX_SPACEDIM));
     102           0 :         for (amrex::MFIter mfi(*a_field[lev], amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
     103             :         {
     104           0 :             amrex::Box bx;
     105           0 :             amrex::IndexType type = a_field[lev]->ixType();
     106           0 :             if (type == amrex::IndexType::TheCellType())      bx = mfi.growntilebox();
     107           0 :             else if (type == amrex::IndexType::TheNodeType()) bx = mfi.grownnodaltilebox();
     108           0 :             else Util::Abort(INFO, "Unkonwn index type");
     109             : 
     110           0 :             Set::Patch<Set::Vector> field = a_field.Patch(lev,mfi);
     111           0 :             for (unsigned int n = 0; n < AMREX_SPACEDIM; n++)
     112             :             {
     113           0 :                 amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k)
     114             :                 {
     115           0 :                     Set::Vector x = Set::Position(i, j, k, geom[lev], type);
     116           0 :                     if (coord == Expression::CoordSys::Cartesian)
     117             :                     {
     118             : #if AMREX_SPACEDIM == 1
     119             :                         field(i, j, k)(n) = f[n](x(0), 0.0, 0.0, a_time);
     120             : #elif AMREX_SPACEDIM == 2
     121           0 :                         field(i, j, k)(n) = f[n](x(0), x(1), 0.0, a_time);
     122             : #elif AMREX_SPACEDIM == 3
     123           0 :                         field(i, j, k)(n) = f[n](x(0), x(1), x(2), a_time);
     124             : #endif
     125             :                     }
     126             : #if AMREX_SPACEDIM>1
     127           0 :                     else if (coord == Expression::CoordSys::Polar)
     128             :                     {
     129           0 :                         field(i, j, k)(n) = f[n](sqrt(x(0)* x(0) + x(1) * x(1)), std::atan2(x(1), x(0)), x(2), a_time);
     130             :                     }
     131             : #endif
     132           0 :                 });
     133             :             }
     134             :         }
     135           0 :         a_field[lev]->FillBoundary();
     136           0 :     };
     137             : 
     138           3 :     static void Parse(Expression& value, IO::ParmParse& pp)
     139             :     {
     140           3 :         for (int i = 0; true; i++)
     141             :         {
     142           6 :             std::string coordstr = "cartesian";
     143             :             // coordinate system to use: "cartesian" (for x,y,z,t) and 
     144             :             // "polar" (for r, theta, z, t)
     145           6 :             pp_query("coord", coordstr); 
     146           6 :             if (coordstr == "cartesian") value.coord = Expression::CoordSys::Cartesian;
     147           0 :             else if (coordstr == "polar") value.coord = Expression::CoordSys::Polar;
     148           0 :             else Util::Abort(INFO, "unsupported coordinates ", coordstr);
     149             : 
     150           6 :             std::string func = "0.0";
     151           6 :             std::string name = "region" + std::to_string(i);
     152             : 
     153           6 :             if (!pp.contains(name.data())) break;
     154           3 :             pp_query(name.data(), func);
     155             : 
     156           3 :             value.parser.push_back(amrex::Parser(func));
     157             : 
     158             :             //
     159             :             // Read in user-defined constants and add them to the parser
     160             :             //
     161           6 :             std::string prefix = pp.getPrefix();
     162           6 :             std::set<std::string> entries = pp.getEntries(prefix + ".constant");//"constant");
     163           3 :             std::set<std::string>::iterator entry;
     164           6 :             for (entry = entries.begin(); entry != entries.end(); entry++)
     165             :             {
     166           6 :                 IO::ParmParse pp;
     167           6 :                 std::string fullname = *entry;
     168           3 :                 Set::Scalar val  = NAN;
     169           3 :                 pp_query(fullname.data(),val);
     170           6 :                 std::string name = Util::String::Split(fullname,'.').back();
     171           3 :                 value.parser.back().setConstant(name,val);
     172             :             }
     173             : 
     174           3 :             if (value.coord == Expression::CoordSys::Cartesian)
     175             :             {
     176          15 :                 value.parser.back().registerVariables({ "x","y","z","t" });
     177           3 :                 value.f.push_back(value.parser.back().compile<4>());
     178             :             }
     179           0 :             else if (value.coord == Expression::CoordSys::Polar)
     180             :             {
     181           0 :                 value.parser.back().registerVariables({ "r","theta","z","t" });
     182           0 :                 value.f.push_back(value.parser.back().compile<4>());
     183             :             }
     184           3 :         }
     185             : 
     186           3 :     };
     187             : };
     188             : }
     189             : 
     190             : #endif