text: /*
   * (C) Copyright 2013 ECMWF.
   *
   * This software is licensed under the terms of the Apache Licence Version 2.0
   * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
   * In applying this licence, ECMWF does not waive the privileges and immunities
   * granted to it by virtue of its status as an intergovernmental organisation
   * nor does it submit to any jurisdiction.
   */
  
  #include <algorithm>
  
  #include "atlas/array.h"
  #include "atlas/field/FieldSet.h"
  #include "atlas/functionspace/NodeColumns.h"
  #include "atlas/functionspace/Spectral.h"
  #include "atlas/functionspace/StructuredColumns.h"
  #include "atlas/grid/Distribution.h"
  #include "atlas/grid/detail/partitioner/EqualRegionsPartitioner.h"
  #include "atlas/grid/detail/partitioner/TransPartitioner.h"
  #include "atlas/mesh/Mesh.h"
  #include "atlas/mesh/Nodes.h"
  #include "atlas/meshgenerator.h"
  #include "atlas/option.h"
  #include "atlas/output/Gmsh.h"
  #include "atlas/parallel/mpi/mpi.h"
  #include "atlas/trans/Trans.h"
  #include "atlas/util/CoordinateEnums.h"
  #include "atlas/util/Earth.h"
  
  #include "tests/AtlasTestEnvironment.h"
  
  #if ATLAS_HAVE_TRANS
  #include "atlas/library/config.h"
  #include "atlas/trans/ifs/TransIFS.h"
  #if ATLAS_HAVE_ECTRANS
  #include "ectrans/transi.h"
  #else
  #include "transi/trans.h"
  #endif
  #endif
  
  namespace atlas {
  namespace test {
  
  //-----------------------------------------------------------------------------
  
  struct AtlasTransEnvironment : public AtlasTestEnvironment {
      AtlasTransEnvironment(int argc, char* argv[]): AtlasTestEnvironment(argc, argv) {
          if (mpi::comm().size() == 1) {
              trans_use_mpi(false);
          }
          trans_init();
      }
  
      ~AtlasTransEnvironment() { trans_finalize(); }
  };
  
  /// @brief Compute magnitude of flow with rotation-angle beta
  /// (beta=0 --> zonal, beta=pi/2 --> meridional)
  static void rotated_flow_magnitude(StructuredGrid& grid, double var[], const double& beta) {
      const double radius  = util::Earth::radius();
      const double USCAL   = 20.;
      const double pvel    = USCAL / radius;
      const double deg2rad = M_PI / 180.;
  
      idx_t n(0);
      for (idx_t jlat = 0; jlat < grid.ny(); ++jlat) {
          for (idx_t jlon = 0; jlon < grid.nx(jlat); ++jlon) {
              const double x = grid.x(jlon, jlat) * deg2rad;
              const double y = grid.y(jlat) * deg2rad;
              const double Ux =
                  pvel * (std::cos(beta) + std::tan(y) * std::cos(x) * std::sin(beta)) * radius * std::cos(y);
              const double Uy = -pvel * std::sin(x) * std::sin(beta) * radius;
              var[n]          = std::sqrt(Ux * Ux + Uy * Uy);
              ++n;
          }
      }
  }
  
  /// @brief Compute magnitude of flow with rotation-angle beta
  /// (beta=0 --> zonal, beta=pi/2 --> meridional)
  void rotated_flow_magnitude(const functionspace::NodeColumns& fs, Field& field, const double& beta) {
      const double radius  = util::Earth::radius();
      const double USCAL   = 20.;
      const double pvel    = USCAL / radius;
      const double deg2rad = M_PI / 180.;
  
      array::ArrayView<double, 2> lonlat_deg = array::make_view<double, 2>(fs.nodes().lonlat());
      array::ArrayView<double, 1> var        = array::make_view<double, 1>(field);
  
      size_t nnodes = fs.nodes().size();
      for (size_t jnode = 0; jnode < nnodes; ++jnode) {
          double x   = lonlat_deg(jnode, (size_t)LON) * deg2rad;
          double y   = lonlat_deg(jnode, (size_t)LAT) * deg2rad;
          double Ux  = pvel * (std::cos(beta) + std::tan(y) * std::cos(x) * std::sin(beta)) * radius * std::cos(y);
          double Uy  = -pvel * std::sin(x) * std::sin(beta) * radius;
          var(jnode) = std::sqrt(Ux * Ux + Uy * Uy);
      }
  }
  
  //-----------------------------------------------------------------------------
  
  CASE("test_invtrans_ifsStyle") {
      std::string grid_uid("O80");
      StructuredGrid g(grid_uid);
      long N = g.ny() / 2;
      trans::TransIFS trans(g, 2 * N - 1);
      Log::info() << "Trans initialized" << std::endl;
      std::vector<double> rspecg;
      int nfld = 1;
  
      std::vector<double> init_gpg(trans.grid().size());
      std::vector<double> init_gp(trans.trans()->ngptot);
      std::vector<double> init_sp(trans.trans()->nspec2);
      std::vector<int> nfrom(nfld, 1);
      if (mpi::comm().rank() == 0) {
          double beta = M_PI * 0.5;
          rotated_flow_magnitude(g, init_gpg.data(), beta);
      }
      trans.distgrid(nfld, nfrom.data(), init_gpg.data(), init_gp.data());
      trans.dirtrans(nfld, init_gp.data(), init_sp.data());
  
      std::vector<double> rgp(3 * nfld * trans.trans()->ngptot);
      double *no_vorticity(nullptr), *no_divergence(nullptr);
      int nb_vordiv(0);
      int nb_scalar(nfld);
      trans.invtrans(nb_scalar, init_sp.data(), nb_vordiv, no_vorticity, no_divergence, rgp.data(),
                     option::scalar_derivatives(true));
  
      std::vector<int> nto(nfld, 1);
      std::vector<double> rgpg(3 * nfld * trans.grid().size());
  
      trans.gathgrid(nfld, nto.data(), rgp.data(), rgpg.data());
  
      // Output
      {
          Mesh mesh = StructuredMeshGenerator().generate(g);
          functionspace::StructuredColumns gp(g);
          output::Gmsh gmsh(grid_uid + "-grid.msh");
          Field scalar("scalar", rgp.data(), array::make_shape(gp.size()));
          Field scalar_dNS("scalar_dNS", rgp.data() + nfld * gp.size(), array::make_shape(gp.size()));
          Field scalar_dEW("scalar_dEW", rgp.data() + 2 * nfld * gp.size(), array::make_shape(gp.size()));
          gmsh.write(mesh);
          gmsh.write(scalar, gp);
          gmsh.write(scalar_dEW, gp);
          gmsh.write(scalar_dNS, gp);
      }
  }
  
  CASE("test_invtrans_grad") {
      std::string grid_uid("O48");
      StructuredGrid g(grid_uid);
      Mesh mesh = StructuredMeshGenerator().generate(g);
      idx_t N   = g.ny() / 2;
      trans::Trans trans(g, 2 * N - 1);
      functionspace::NodeColumns gp(mesh);
      functionspace::Spectral sp(trans);
  
      Field scalar_sp = sp.createField<double>(option::name("scalar_sp"));
      Field scalar    = gp.createField<double>(option::name("scalar"));
      Field grad      = gp.createField<double>(option::name("grad") | option::variables(2));
  
      // Initial condition
      double beta = M_PI * 0.5;
      rotated_flow_magnitude(gp, scalar, beta);
  
      // Transform to spectral
      trans.dirtrans(scalar, scalar_sp);
  
      // Inverse transform for gradient
      trans.invtrans_grad(scalar_sp, grad);
  
      gp.haloExchange(grad);
  
      // Output
      {
          Mesh mesh = StructuredMeshGenerator().generate(g);
          functionspace::StructuredColumns gp(g);
          output::Gmsh gmsh(grid_uid + "-nodes.msh");
          gmsh.write(mesh);
          gmsh.write(scalar, gp);
          gmsh.write(grad, gp);
      }
  }
  
  //-----------------------------------------------------------------------------
  
  }  // namespace test
  }  // namespace atlas
  
  int main(int argc, char** argv) {
      return atlas::test::run<atlas::test::AtlasTransEnvironment>(argc, argv);
  }
  