2
2
#define EAMXX_MAM_FIND_SEASON_INDEX_UTILS
3
3
4
4
#include < ekat/kokkos/ekat_kokkos_utils.hpp>
5
+ #include < mam4xx/mam4.hpp>
6
+
5
7
#include " share/io/scorpio_input.hpp"
6
8
#include " share/io/scream_scorpio_interface.hpp"
7
- #include < mam4xx/mam4.hpp>
8
9
9
10
namespace scream ::mam_coupling {
10
11
@@ -13,27 +14,28 @@ using ESU = ekat::ExeSpaceUtils<ExeSpace>;
13
14
14
15
// views for single- and multi-column data
15
16
16
- using const_view_1d = typename KT::template view_1d<const Real>;
17
- using view_int_2d = typename KT::template view_2d<int >;
17
+ using const_view_1d = typename KT::template view_1d<const Real>;
18
+ using view_int_2d = typename KT::template view_2d<int >;
18
19
19
- using view_1d_host = typename KT::view_1d<Real>::HostMirror;
20
- using view_3d_host = typename KT::view_3d<Real>::HostMirror;
20
+ using view_1d_host = typename KT::view_1d<Real>::HostMirror;
21
+ using view_3d_host = typename KT::view_3d<Real>::HostMirror;
21
22
using view_int_3d_host = typename KT::view_3d<int >::HostMirror;
22
23
using view_int_2d_host = typename KT::view_2d<int >::HostMirror;
23
24
24
25
/* *
25
- * @brief Reads the season index from the given file and computes the season indices based on latitudes.
26
+ * @brief Reads the season index from the given file and computes the season
27
+ * indices based on latitudes.
26
28
*
27
29
* @param[in] season_wes_file The path to the season_wes.nc file.
28
30
* @param[in] clat A 1D view of latitude values in degrees.
29
- * @param[out] index_season_lai A 2D view to store the computed season indices. Note that indices are in C++ (starting from zero).
31
+ * @param[out] index_season_lai A 2D view to store the computed season indices.
32
+ * Note that indices are in C++ (starting from zero).
30
33
*/
31
34
32
- inline void find_season_index_reader (const std::string& season_wes_file,
33
- const const_view_1d& clat,
34
- view_int_2d &index_season_lai)
35
- {
36
- const int plon= clat.extent (0 );
35
+ inline void find_season_index_reader (const std::string &season_wes_file,
36
+ const const_view_1d &clat,
37
+ view_int_2d &index_season_lai) {
38
+ const int plon = clat.extent (0 );
37
39
scorpio::register_file (season_wes_file, scorpio::Read);
38
40
39
41
const int nlat_lai = scorpio::get_dimlen (season_wes_file, " lat" );
@@ -46,34 +48,32 @@ inline void find_season_index_reader(const std::string& season_wes_file,
46
48
scorpio::read_var (season_wes_file, " lat" , lat_lai.data ());
47
49
48
50
Kokkos::MDRangePolicy<Kokkos::HostSpace::execution_space, Kokkos::Rank<2 >>
49
- policy_wk_lai ({0 , 0 }, {nlat_lai, npft_lai});
51
+ policy_wk_lai ({0 , 0 }, {nlat_lai, npft_lai});
50
52
51
53
// loop over time to get all 12 instantence of season_wes
52
- for (int itime = 0 ; itime < 12 ; ++itime)
53
- {
54
- scorpio::read_var (season_wes_file, " season_wes" , wk_lai_temp.data (),itime);
55
- // copy data from wk_lai_temp to wk_lai.
56
- // NOTE: season_wes has different layout that wk_lai
57
- Kokkos::parallel_for (" copy_to_wk_lai" , policy_wk_lai,
58
- [&](const int j, const int k) {
59
- wk_lai (j, k, itime) = wk_lai_temp (k, j);
60
- });
61
- Kokkos::fence ();
54
+ for (int itime = 0 ; itime < 12 ; ++itime) {
55
+ scorpio::read_var (season_wes_file, " season_wes" , wk_lai_temp.data (), itime);
56
+ // copy data from wk_lai_temp to wk_lai.
57
+ // NOTE: season_wes has different layout that wk_lai
58
+ Kokkos::parallel_for (" copy_to_wk_lai" , policy_wk_lai,
59
+ [&](const int j, const int k) {
60
+ wk_lai (j, k, itime) = wk_lai_temp (k, j);
61
+ });
62
+ Kokkos::fence ();
62
63
}
63
64
scorpio::release_file (season_wes_file);
64
65
65
- index_season_lai =view_int_2d (" index_season_lai" , plon, 12 );
66
- const view_int_2d_host index_season_lai_host= Kokkos::create_mirror_view (index_season_lai);
66
+ index_season_lai = view_int_2d (" index_season_lai" , plon, 12 );
67
+ const view_int_2d_host index_season_lai_host =
68
+ Kokkos::create_mirror_view (index_season_lai);
67
69
68
- const view_1d_host clat_host= Kokkos::create_mirror_view (clat);
70
+ const view_1d_host clat_host = Kokkos::create_mirror_view (clat);
69
71
Kokkos::deep_copy (clat_host, clat);
70
72
71
73
// Computation is performed on the host
72
74
mam4::mo_drydep::find_season_index (clat_host, lat_lai, nlat_lai, wk_lai,
73
- index_season_lai_host);
75
+ index_season_lai_host);
74
76
Kokkos::deep_copy (index_season_lai, index_season_lai_host);
75
-
76
-
77
77
}
78
78
} // namespace scream::mam_coupling
79
79
#endif //
0 commit comments