Skip to content

Commit

Permalink
Merge branch 'pass_fields_to_mixers' into 'master'
Browse files Browse the repository at this point in the history 
Pass fields to mixers instead of arrays

See merge request npneq/inq!1129
  • Loading branch information
@xavierandrade
xavierandrade committed Sep 20, 2024
2 parents e23534d + 1b37e01 commit 46577d6
Show file tree
Hide file tree
Showing 5 changed files with 78 additions and 57 deletions.
8 changes: 3 additions & 5 deletions src/ground_state/calculator.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,12 +121,12 @@ class calculator {
results res;
operations::preconditioner prec;

using mix_arr_type = std::remove_reference_t<decltype(electrons.spin_density().matrix().flatted())>;
using mix_arr_type = std::remove_reference_t<decltype(electrons.spin_density())>;

auto mixer = [&]()->std::unique_ptr<mixers::base<mix_arr_type>>{
switch(solver_.mixing_algorithm()){
case options::ground_state::mixing_algo::LINEAR : return std::make_unique<mixers::linear <mix_arr_type>>(solver_.mixing());
case options::ground_state::mixing_algo::BROYDEN: return std::make_unique<mixers::broyden<mix_arr_type>>(4, solver_.mixing(), electrons.spin_density().matrix().flatted().size(), electrons.density_basis().comm());
case options::ground_state::mixing_algo::BROYDEN: return std::make_unique<mixers::broyden<mix_arr_type>>(4, solver_.mixing(), electrons.spin_density().matrix().flatted().size());
} __builtin_unreachable();
}();

Expand Down Expand Up @@ -191,9 +191,7 @@ class calculator {
density_diff /= electrons.states().num_electrons();

if(inter_.self_consistent()) {
auto tmp = +electrons.spin_density().matrix().flatted();
mixer->operator()(tmp, new_density.matrix().flatted());
electrons.spin_density().matrix().flatted() = tmp;
mixer->operator()(electrons.spin_density(), new_density);
observables::density::normalize(electrons.spin_density(), electrons.states().num_electrons());
} else {
electrons.spin_density() = std::move(new_density);
Expand Down
4 changes: 2 additions & 2 deletions src/mixers/base.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,12 @@
namespace inq {
namespace mixers {

template <class ArrayType>
template <class FieldType>
class base {

public:
virtual ~base(){};
virtual void operator()(ArrayType & input_value, ArrayType const & output_value) = 0;
virtual void operator()(FieldType & input_value, FieldType const & output_value) = 0;

};

Expand Down
74 changes: 41 additions & 33 deletions src/mixers/broyden.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,29 +22,27 @@
namespace inq {
namespace mixers {

template <class ArrayType>
class broyden : public base<ArrayType> {
template <class FieldType>
class broyden : public base<FieldType> {

public:

using element_type = typename ArrayType::element_type;
using element_type = typename FieldType::element_type;

template <class CommType>
broyden(const int arg_steps, const double arg_mix_factor, const long long dim, CommType & comm):
broyden(const int arg_steps, const double arg_mix_factor, const long long dim):
iter_(0),
max_size_(arg_steps),
mix_factor_(arg_mix_factor),
dv_({max_size_, dim}, NAN),
df_({max_size_, dim}, NAN),
f_old_(dim, NAN),
vin_old_(dim, NAN),
last_pos_(-1),
comm_(comm){
last_pos_(-1){
}

///////////////////////////////////////////////////////////////////////////////////////////////////
void broyden_extrapolation(ArrayType & input_value, int const iter_used, gpu::array<element_type, 1> const & ff){
template <typename Comm, typename Array>
void broyden_extrapolation(Comm & comm, Array & input_value, int const iter_used, gpu::array<element_type, 1> const & ff){

CALI_CXX_MARK_SCOPE("broyden_extrapolation");

Expand All @@ -70,12 +68,12 @@ class broyden : public base<ArrayType> {
auto workmat = +blas::gemm(1.0, matff, blas::H(subdf));
auto work = +workmat[0];

gpu::run(iter_used, [w0, ww, be = begin(beta), dfactor = 1.0/comm_.size()] GPU_LAMBDA (auto ii){ be[ii][ii] = dfactor*(w0*w0 + ww*ww); });
gpu::run(iter_used, [w0, ww, be = begin(beta), dfactor = 1.0/comm.size()] GPU_LAMBDA (auto ii){ be[ii][ii] = dfactor*(w0*w0 + ww*ww); });

if(comm_.size() > 1){
if(comm.size() > 1){
CALI_CXX_MARK_SCOPE("broyden_extrapolation::reduce");
comm_.all_reduce_n(raw_pointer_cast(beta.data_elements()), beta.num_elements());
comm_.all_reduce_n(raw_pointer_cast(work.data_elements()), work.num_elements());
comm.all_reduce_n(raw_pointer_cast(beta.data_elements()), beta.num_elements());
comm.all_reduce_n(raw_pointer_cast(work.data_elements()), work.num_elements());
}

solvers::least_squares(beta, work);
Expand All @@ -92,10 +90,13 @@ class broyden : public base<ArrayType> {

///////////////////////////////////////////////////////////////////////////////////////////////////

void operator()(ArrayType & input_value, ArrayType const & output_value){
void operator()(FieldType & input_field, FieldType const & output_field){

CALI_CXX_MARK_SCOPE("broyden_mixing");

auto input_value = input_field.matrix().flatted();
auto output_value = output_field.matrix().flatted();

assert((typename gpu::array<double, 2>::size_type) input_value.size() == dv_[0].size());
assert((typename gpu::array<double, 2>::size_type) output_value.size() == dv_[0].size());

Expand Down Expand Up @@ -126,9 +127,9 @@ class broyden : public base<ArrayType> {

gamma_ = dot(conj(df_[pos]), df_[pos]);

if(comm_.size() > 1){
if(input_field.full_comm().size() > 1){
CALI_CXX_MARK_SCOPE("broyden_mixing::reduce");
comm_.all_reduce_in_place_n(&gamma_, 1, std::plus<>{});
input_field.full_comm().all_reduce_in_place_n(&gamma_, 1, std::plus<>{});
}

gamma_ = std::max(1e-8, sqrt(gamma_));
Expand All @@ -148,7 +149,7 @@ class broyden : public base<ArrayType> {

auto iter_used = std::min(iter_ - 1, max_size_);

broyden_extrapolation(input_value, iter_used, ff);
broyden_extrapolation(input_field.full_comm(), input_value, iter_used, ff);

}

Expand All @@ -163,8 +164,6 @@ class broyden : public base<ArrayType> {
gpu::array<element_type, 1> vin_old_;
element_type gamma_;
int last_pos_;
mutable parallel::communicator comm_;

};

}
Expand All @@ -180,24 +179,33 @@ TEST_CASE(INQ_TEST_FILE, INQ_TEST_TAG) {

using namespace inq;
using namespace Catch::literals;


gpu::array<double, 1> vin({10.0, -20.0});
gpu::array<double, 1> vout({0.0, 22.2});

mixers::broyden<decltype(vin)> lm(5, 0.5, 2, boost::mpi3::environment::get_self_instance());
auto comm = parallel::communicator{boost::mpi3::environment::get_self_instance()};

lm(vin, vout);

CHECK(vin[0] == 5.0_a);
CHECK(vin[1] == 1.1_a);

vout = gpu::array<double, 1>({4.0, 5.5});
basis::trivial bas(2, comm);

basis::field_set<basis::trivial, double> vin(bas, 1);
vin.matrix()[0][0] = 10.0;
vin.matrix()[1][0] = -20.0;

lm(vin, vout);
basis::field_set<basis::trivial, double> vout(bas, 1);
vout.matrix()[0][0] = 0.0;
vout.matrix()[1][0] = 22.2;

CHECK(vin[0] == 4.4419411001_a);
CHECK(vin[1] == 3.5554591594_a);
mixers::broyden<decltype(vin)> mixer(5, 0.5, 2);

mixer(vin, vout);

CHECK(vin.matrix()[0][0] == 5.0_a);
CHECK(vin.matrix()[1][0] == 1.1_a);

vout.matrix()[0][0] = 4.0;
vout.matrix()[1][0] = 5.5;

mixer(vin, vout);

CHECK(vin.matrix()[0][0] == 4.4419411001_a);
CHECK(vin.matrix()[1][0] == 3.5554591594_a);

}
#endif
41 changes: 29 additions & 12 deletions src/mixers/linear.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,24 +17,28 @@
namespace inq {
namespace mixers {

template <class ArrayType>
class linear : public base<ArrayType> {
template <class FieldType>
class linear : public base<FieldType> {

public:

linear(double arg_mix_factor):
mix_factor_(arg_mix_factor){
}

void operator()(ArrayType & input_value, ArrayType const & output_value){
void operator()(FieldType & input_value, FieldType const & output_value){
//note: arguments might alias here

assert(input_value.basis() == output_value.basis());
assert(input_value.local_set_size() == output_value.local_set_size());

CALI_CXX_MARK_SCOPE("linear_mixing");

gpu::run(input_value.size(),
[iv = begin(input_value), ov = begin(output_value), mix = mix_factor_] GPU_LAMBDA (auto ii){
iv[ii] = (1.0 - mix)*iv[ii] + mix*ov[ii];
gpu::run(input_value.local_set_size(), input_value.basis().size(),
[iv = begin(input_value.matrix()), ov = begin(output_value.matrix()), mix = mix_factor_] GPU_LAMBDA (auto is, auto ip){
iv[ip][is] = (1.0 - mix)*iv[ip][is] + mix*ov[ip][is];
});

}

private:
Expand All @@ -57,15 +61,28 @@ TEST_CASE(INQ_TEST_FILE, INQ_TEST_TAG) {
using namespace inq;
using namespace Catch::literals;

gpu::array<double, 1> vin({10.0, -20.0});
gpu::array<double, 1> vout({0.0, 22.2});
basis::trivial bas(2, parallel::communicator{boost::mpi3::environment::get_self_instance()});

basis::field_set<basis::trivial, double> vin(bas, 2);
vin.matrix()[0][0] = 10.0;
vin.matrix()[1][0] = -20.0;
vin.matrix()[0][1] = 3.45;
vin.matrix()[1][1] = 192.34;

basis::field_set<basis::trivial, double> vout(bas, 2);
vout.matrix()[0][0] = 0.0;
vout.matrix()[1][0] = 22.2;
vout.matrix()[0][1] = 1.87;
vout.matrix()[1][1] = -133.55;

mixers::linear<decltype(vin)> lm(0.5);
mixers::linear<decltype(vin)> mixer(0.5);

lm(vin, vout);
mixer(vin, vout);

CHECK(vin[0] == 5.0_a);
CHECK(vin[1] == 1.1_a);
CHECK(vin.matrix()[0][0] == 5.0_a);
CHECK(vin.matrix()[1][0] == 1.1_a);
CHECK(vin.matrix()[0][1] == 2.66_a);
CHECK(vin.matrix()[1][1] == 29.395_a);

}
#endif
8 changes: 3 additions & 5 deletions src/real_time/crank_nicolson.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,8 @@ void crank_nicolson(double const time, double const dt, systems::ions & ions, sy
energy.ion(ionic::interaction_energy(ions.cell(), ions, electrons.atomic_pot()));
}

using mix_arr_type = std::remove_reference_t<decltype(electrons.spin_density().matrix().flatted())>;
auto mixer = mixers::broyden<mix_arr_type>(4, 0.3, electrons.spin_density().matrix().flatted().size(), electrons.density_basis().comm());
using mix_arr_type = std::remove_reference_t<decltype(electrons.spin_density())>;
auto mixer = mixers::broyden<mix_arr_type>(4, 0.3, electrons.spin_density().matrix().flatted().size());

auto old_exxe = 0.0;
auto update_hf = true;
Expand Down Expand Up @@ -103,9 +103,7 @@ void crank_nicolson(double const time, double const dt, systems::ions & ions, sy
if(exxe_diff < exxe_tol) break;
}

auto tmp = +electrons.spin_density().matrix().flatted();
mixer(tmp, new_density.matrix().flatted());
electrons.spin_density().matrix().flatted() = tmp;
mixer(electrons.spin_density(), new_density);
observables::density::normalize(electrons.spin_density(), electrons.states().num_electrons());
}

Expand Down

0 comments on commit 46577d6

Please sign in to comment.