diff --git a/src/loss_function/LossFunctionLogPearsonRMSE.cpp b/src/loss_function/LossFunctionLogPearsonRMSE.cpp
index bb7c63a937e60f508bd5e767172be202e2e128ef..422f6f0532287d959f9f547e9ffb8cb6d5fdf28f 100644
--- a/src/loss_function/LossFunctionLogPearsonRMSE.cpp
+++ b/src/loss_function/LossFunctionLogPearsonRMSE.cpp
@@ -64,19 +64,19 @@ void LossFunctionLogPearsonRMSE::set_a(const std::vector<int>& inds, int taskind
}
}
-void LossFunctionLogPearsonRMSE::set_a(const std::vector<std::vector<int>>& feature_indices,
+void LossFunctionLogPearsonRMSE::set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models,
int taskind,
int start)
{
- for (int model_idx = 0; model_idx < feature_indices.size(); ++model_idx)
+ for (int model_idx = 0; model_idx < models.extent(1); ++model_idx)
{
- for (int ff = 0; ff < feature_indices[model_idx].size(); ++ff)
+ for (int feature_idx = 0; feature_idx < models.extent(0); ++feature_idx)
{
- double* val_ptr = node_value_arrs::get_d_matrix_ptr(feature_indices[model_idx][ff]) +
+ double* val_ptr = node_value_arrs::get_d_matrix_ptr(models(feature_idx, model_idx)) +
start;
std::transform(val_ptr,
val_ptr + _task_sizes_train[taskind],
- &_a(0, ff, model_idx),
+ &_a(0, feature_idx, model_idx),
[](double val) { return std::log(val); });
}
}
diff --git a/src/loss_function/LossFunctionLogPearsonRMSE.hpp b/src/loss_function/LossFunctionLogPearsonRMSE.hpp
index 29232992ee929cbf0344ff5e14561e7585a1f0a8..3d85c1a5ccacc8cadbf8c91f27d1bfa522427e15 100644
--- a/src/loss_function/LossFunctionLogPearsonRMSE.hpp
+++ b/src/loss_function/LossFunctionLogPearsonRMSE.hpp
@@ -97,7 +97,7 @@ public:
*/
void set_a(const std::vector<int>& inds, int taskind, int start);
- void set_a(const std::vector<std::vector<int>>& feature_indices, int taskind, int start) override;
+ void set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models, int taskind, int start) override;
/**
* @brief Set the A matrix used for solving the least squares regression
diff --git a/src/loss_function/LossFunctionPearsonRMSE.cpp b/src/loss_function/LossFunctionPearsonRMSE.cpp
index 07e8d34e036ec5f5b8495c3951f63a8b97c9c699..ebea62fa85f171b1c824a9d190985328562df7e2 100644
--- a/src/loss_function/LossFunctionPearsonRMSE.cpp
+++ b/src/loss_function/LossFunctionPearsonRMSE.cpp
@@ -86,6 +86,7 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(std::vector<double> prop_train,
_c("c", _n_task),
_training_properties("batched_scores", 0),
_batched_scores("batched_scores"),
+ _models("models", _n_feat),
_lwork(0),
_estimated_training_properties("estimated_training_properties", 0)
{
@@ -114,6 +115,7 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(std::shared_ptr<LossFunction> o
_c("c", _n_task),
_training_properties("batched_scores", 0),
_batched_scores("batched_scores"),
+ _models("models", _n_feat),
_lwork(0),
_estimated_training_properties("estimated_training_properties", 0)
{
@@ -142,6 +144,7 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(const LossFunctionPearsonRMSE&
_c("c", _n_task),
_training_properties("batched_scores", 0),
_batched_scores("batched_scores"),
+ _models("models", _n_feat),
_lwork(0),
_estimated_training_properties("estimated_training_properties", 0)
{
@@ -218,6 +221,8 @@ void LossFunctionPearsonRMSE::set_nfeat(int n_feat)
{
_training_properties(material_idx) = _prop_train[material_idx];
}
+
+ Kokkos::resize(_models, _n_feat);
}
void LossFunctionPearsonRMSE::reset_projection_prop(
@@ -324,17 +329,27 @@ Kokkos::View<double*> LossFunctionPearsonRMSE::operator()(
const size_t batch_size = feature_indices.size();
assert(batch_size <= MAX_BATCHES);
+ assert(feature_indices.size() <= _models.extent(1));
+ assert(feature_indices[0].size() == _models.extent(0));
+ for (int model_idx = 0; model_idx < feature_indices.size(); ++model_idx)
+ {
+ for (int ff = 0; ff < feature_indices[model_idx].size(); ++ff)
+ {
+ _models(ff, model_idx) = feature_indices[model_idx][ff];
+ }
+ }
+
int start = 0;
for (int task_idx = 0; task_idx < _n_task; ++task_idx)
{
Kokkos::deep_copy(_a, 1.0);
- set_a(feature_indices, task_idx, start);
+ set_a(_models, task_idx, start);
set_b(task_idx, start);
Kokkos::fence();
least_squares(task_idx, start);
Kokkos::deep_copy(_a, 1.0);
- set_a(feature_indices, task_idx, start);
+ set_a(_models, task_idx, start);
Kokkos::fence();
set_prop_train_est(_estimated_training_properties, task_idx, start);
@@ -392,36 +407,25 @@ double LossFunctionPearsonRMSE::test_loss(const std::vector<model_node_ptr>& fea
void LossFunctionPearsonRMSE::set_a(const std::vector<int>& inds, int taskind, int start)
{
- set_a(std::vector<std::vector<int>>(1, inds), taskind, start);
+ for (int ff = 0; ff < inds.size(); ++ff)
+ {
+ _models(ff, 0) = inds[ff];
+ }
+ set_a(_models, taskind, start);
}
-void LossFunctionPearsonRMSE::set_a(const std::vector<std::vector<int>>& feature_indices,
+void LossFunctionPearsonRMSE::set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models,
int taskind,
int start)
{
- Kokkos::View<int**, Kokkos::LayoutLeft> models(
- "models", feature_indices[0].size(), feature_indices.size());
- auto max_idx = 0;
- for (int model_idx = 0; model_idx < feature_indices.size(); ++model_idx)
- {
- for (int ff = 0; ff < feature_indices[model_idx].size(); ++ff)
- {
- models(ff, model_idx) = feature_indices[model_idx][ff];
- max_idx = std::max(max_idx, feature_indices[model_idx][ff]);
- }
- }
-
assert(_descriptor_matrix.extent(0) >= _task_sizes_train[taskind] + start);
- assert(_descriptor_matrix.extent(1) > max_idx);
Kokkos::deep_copy(_a, 1.0);
auto a = _a;
auto descriptor_matrix = _descriptor_matrix;
auto policy = Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
{0, 0, 0},
- {static_cast<size_t>(_task_sizes_train[taskind]),
- feature_indices[0].size(),
- feature_indices.size()});
+ {static_cast<size_t>(_task_sizes_train[taskind]), models.extent(0), models.extent(1)});
auto kernel = KOKKOS_LAMBDA(const int sample_idx, const int feature_idx, const int model_idx)
{
a(sample_idx, feature_idx, model_idx) = descriptor_matrix(sample_idx + start,
diff --git a/src/loss_function/LossFunctionPearsonRMSE.hpp b/src/loss_function/LossFunctionPearsonRMSE.hpp
index 246fcd8c9be2d1cb1685d644be6186af87aeb5cb..0b7757aab6928f2a645e7b5488a0913c362c4b6e 100644
--- a/src/loss_function/LossFunctionPearsonRMSE.hpp
+++ b/src/loss_function/LossFunctionPearsonRMSE.hpp
@@ -56,6 +56,7 @@ protected:
Kokkos::View<double*> _work; //!< Work vector for dgels
Kokkos::View<double*> _training_properties;
Kokkos::View<double[MAX_BATCHES]> _batched_scores;
+ Kokkos::View<int*[MAX_BATCHES], Kokkos::LayoutLeft> _models;
Kokkos::View<double**, Kokkos::LayoutLeft> _descriptor_matrix;
Kokkos::View<double* [MAX_BATCHES], Kokkos::LayoutLeft> _estimated_training_properties;
@@ -175,7 +176,7 @@ public:
*/
virtual void set_a(const std::vector<int>& inds, int taskind, int start);
- virtual void set_a(const std::vector<std::vector<int>>& feature_indices, int taskind, int start);
+ virtual void set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models, int taskind, int start);
/**
* @brief Set the A matrix used for solving the least squares regression