Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reorder Jacobian calculation #707

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from
+160 −107
Draft

Reorder Jacobian calculation #707

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
986ec3a
calculate jacobian by column
mattldawson Dec 26, 2024
68250d7
Merge branch 'main' into develop-reorder-jacobian
mattldawson Jan 9, 2025
e7effcb
Merge branch 'main' into develop-reorder-jacobian
mattldawson Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
201 changes: 121 additions & 80 deletions include/micm/process/process_set.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,11 +60,24 @@ namespace micm
class ProcessSet
{
protected:
/// @brief Process information for use in setting Jacobian elements
struct ProcessInfo
{
std::size_t process_id_;
std::size_t independent_id_;
std::size_t number_of_dependent_reactants_;
std::size_t number_of_products_;
};

std::vector<std::size_t> number_of_reactants_;
std::vector<std::size_t> reactant_ids_;
std::vector<std::size_t> number_of_products_;
std::vector<std::size_t> product_ids_;
std::vector<double> yields_;
std::vector<ProcessInfo> jacobian_process_info_;
std::vector<std::size_t> jacobian_reactant_ids_;
std::vector<std::size_t> jacobian_product_ids_;
std::vector<double> jacobian_yields_;
std::vector<std::size_t> jacobian_flat_ids_;

public:
Expand All @@ -82,7 +95,9 @@ namespace micm
/// @return Jacobian elements as a set of index pairs
std::set<std::pair<std::size_t, std::size_t>> NonZeroJacobianElements() const;

/// @brief Sets the indicies for each non-zero Jacobian element in the underlying vector
/// @brief Sets the indicies for each non-zero Jacobian element in the underlying vector.
/// Also sets combination of process ids and reactant ids to allow setting of
/// jacobian elements in order by column.
/// @param matrix The sparse matrix used for the Jacobian
template<typename OrderingPolicy>
void SetJacobianFlatIds(const SparseMatrix<double, OrderingPolicy>& matrix);
Expand Down Expand Up @@ -135,15 +150,21 @@ namespace micm
reactant_ids_(),
number_of_products_(),
product_ids_(),
yields_()
yields_(),
jacobian_process_info_(),
jacobian_reactant_ids_(),
jacobian_product_ids_(),
jacobian_yields_(),
jacobian_flat_ids_()
{
MICM_PROFILE_FUNCTION();

for (auto& process : processes)
// Set up process information for forcing calculations
for (const auto& process : processes)
{
std::size_t number_of_reactants = 0;
std::size_t number_of_products = 0;
for (auto& reactant : process.reactants_)
for (const auto& reactant : process.reactants_)
{
if (reactant.IsParameterized())
continue; // Skip reactants that are parameterizations
Expand All @@ -152,7 +173,7 @@ namespace micm
reactant_ids_.push_back(variable_map.at(reactant.name_));
++number_of_reactants;
}
for (auto& product : process.products_)
for (const auto& product : process.products_)
{
if (product.first.IsParameterized())
continue; // Skip products that are parameterizations
Expand All @@ -165,6 +186,58 @@ namespace micm
number_of_reactants_.push_back(number_of_reactants);
number_of_products_.push_back(number_of_products);
}

// Set up process information for Jacobian calculations
std::vector<std::pair<std::string, std::size_t>> sorted_names(variable_map.begin(), variable_map.end());
std::sort(sorted_names.begin(), sorted_names.end(), [](const auto& a, const auto& b) {
return a.second < b.second;
});
for (const auto& independent_variable : sorted_names)
{
for (std::size_t i_process = 0; i_process < processes.size(); ++i_process)
{
const auto& process = processes[i_process];
// Look for processes that depend on the independent variable
bool found = false;
for (const auto& reactant : process.reactants_)
{
if (reactant.name_ == independent_variable.first)
{
found = true;
break;
}
}
if (!found)
continue;
ProcessInfo info;
info.process_id_ = i_process;
info.independent_id_ = independent_variable.second;
info.number_of_dependent_reactants_ = 0;
info.number_of_products_ = 0;
for (const auto& reactant : process.reactants_)
{
if (reactant.IsParameterized())
continue; // Skip reactants that are parameterizations
if (variable_map.count(reactant.name_) < 1)
throw std::system_error(make_error_code(MicmProcessSetErrc::ReactantDoesNotExist), reactant.name_);
if (variable_map.at(reactant.name_) == independent_variable.second) // skip the independent variable
continue;
jacobian_reactant_ids_.push_back(variable_map.at(reactant.name_));
++info.number_of_dependent_reactants_;
}
for (const auto& product : process.products_)
{
if (product.first.IsParameterized())
continue; // Skip products that are parameterizations
if (variable_map.count(product.first.name_) < 1)
throw std::system_error(make_error_code(MicmProcessSetErrc::ProductDoesNotExist), product.first.name_);
jacobian_product_ids_.push_back(variable_map.at(product.first.name_));
jacobian_yields_.push_back(product.second);
++info.number_of_products_;
}
jacobian_process_info_.push_back(info);
}
}
};

inline std::set<std::pair<std::size_t, std::size_t>> ProcessSet::NonZeroJacobianElements() const
Expand Down Expand Up @@ -199,23 +272,15 @@ namespace micm
MICM_PROFILE_FUNCTION();

jacobian_flat_ids_.clear();
auto react_id = reactant_ids_.begin();
auto prod_id = product_ids_.begin();
for (std::size_t i_rxn = 0; i_rxn < number_of_reactants_.size(); ++i_rxn)
auto react_id = jacobian_reactant_ids_.begin();
auto prod_id = jacobian_product_ids_.begin();
for (const auto& process_info : jacobian_process_info_)
{
for (std::size_t i_ind = 0; i_ind < number_of_reactants_[i_rxn]; ++i_ind)
{
for (std::size_t i_dep = 0; i_dep < number_of_reactants_[i_rxn]; ++i_dep)
{
jacobian_flat_ids_.push_back(matrix.VectorIndex(0, react_id[i_dep], react_id[i_ind]));
}
for (std::size_t i_dep = 0; i_dep < number_of_products_[i_rxn]; ++i_dep)
{
jacobian_flat_ids_.push_back(matrix.VectorIndex(0, prod_id[i_dep], react_id[i_ind]));
}
}
react_id += number_of_reactants_[i_rxn];
prod_id += number_of_products_[i_rxn];
for (std::size_t i_dep = 0; i_dep < process_info.number_of_dependent_reactants_; ++i_dep)
jacobian_flat_ids_.push_back(matrix.VectorIndex(0, *(react_id++), process_info.independent_id_));
jacobian_flat_ids_.push_back(matrix.VectorIndex(0, process_info.independent_id_, process_info.independent_id_));
for (std::size_t i_dep = 0; i_dep < process_info.number_of_products_; ++i_dep)
jacobian_flat_ids_.push_back(matrix.VectorIndex(0, *(prod_id++), process_info.independent_id_));
}
}

Expand Down Expand Up @@ -326,39 +391,22 @@ namespace micm
auto cell_rate_constants = rate_constants[i_cell];
auto cell_state = state_variables[i_cell];

auto react_id = reactant_ids_.begin();
auto yield = yields_.begin();
auto react_id = jacobian_reactant_ids_.begin();
auto yield = jacobian_yields_.begin();
auto flat_id = jacobian_flat_ids_.begin();

// loop over reactions
for (std::size_t i_rxn = 0; i_rxn < number_of_reactants_.size(); ++i_rxn)
// loop over process-dependent variable pairs
for (const auto& process_info : jacobian_process_info_)
{
// loop over number of reactants of a reaction
for (std::size_t i_ind = 0; i_ind < number_of_reactants_[i_rxn]; ++i_ind)
{
double d_rate_d_ind = cell_rate_constants[i_rxn];

for (std::size_t i_react = 0; i_react < number_of_reactants_[i_rxn]; ++i_react)
{
if (i_react == i_ind)
{
continue;
}
d_rate_d_ind *= cell_state[react_id[i_react]];
}
for (std::size_t i_dep = 0; i_dep < number_of_reactants_[i_rxn]; ++i_dep)
{
cell_jacobian[*(flat_id++)] += d_rate_d_ind;
}
for (std::size_t i_dep = 0; i_dep < number_of_products_[i_rxn]; ++i_dep)
{
cell_jacobian[*(flat_id++)] -= yield[i_dep] * d_rate_d_ind;
}
}
react_id += number_of_reactants_[i_rxn];
yield += number_of_products_[i_rxn];
double d_rate_d_ind = cell_rate_constants[process_info.process_id_];
for (std::size_t i_react = 0; i_react < process_info.number_of_dependent_reactants_; ++i_react)
d_rate_d_ind *= cell_state[*(react_id++)];
for (std::size_t i_dep = 0; i_dep < process_info.number_of_dependent_reactants_+1; ++i_dep)
cell_jacobian[*(flat_id++)] += d_rate_d_ind;
for (std::size_t i_dep = 0; i_dep < process_info.number_of_products_; ++i_dep)
cell_jacobian[*(flat_id++)] -= *(yield++) * d_rate_d_ind;
}
// increment cell_jacobian after each row
// increment cell_jacobian after each grid cell
cell_jacobian += jacobian.FlatBlockSize();
}
}
Expand All @@ -382,44 +430,37 @@ namespace micm
// loop over all rows
for (std::size_t i_group = 0; i_group < state_variables.NumberOfGroups(); ++i_group)
{
auto react_id = reactant_ids_.begin();
auto yield = yields_.begin();
auto react_id = jacobian_reactant_ids_.begin();
auto yield = jacobian_yields_.begin();
const std::size_t offset_rc = i_group * rate_constants.GroupSize();
const std::size_t offset_state = i_group * state_variables.GroupSize();
const std::size_t offset_jacobian = i_group * jacobian.GroupSize();
auto flat_id = jacobian_flat_ids_.begin();

for (std::size_t i_rxn = 0; i_rxn < number_of_reactants_.size(); ++i_rxn)
for (const auto& process_info : jacobian_process_info_)
{
for (std::size_t i_ind = 0; i_ind < number_of_reactants_[i_rxn]; ++i_ind)
auto v_rate_subrange_begin = v_rate_constants_begin + offset_rc + (process_info.process_id_ * L);
d_rate_d_ind.assign(v_rate_subrange_begin, v_rate_subrange_begin + L);
for (std::size_t i_react = 0; i_react < process_info.number_of_dependent_reactants_; ++i_react)
{
auto v_rate_subrange_begin = v_rate_constants_begin + offset_rc + (i_rxn * L);
d_rate_d_ind.assign(v_rate_subrange_begin, v_rate_subrange_begin + L);
for (std::size_t i_react = 0; i_react < number_of_reactants_[i_rxn]; ++i_react)
{
if (i_react == i_ind)
{
continue;
}
const std::size_t idx_state_variables = offset_state + (react_id[i_react] * L);
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
d_rate_d_ind[i_cell] *= v_state_variables[idx_state_variables + i_cell];
}
for (std::size_t i_dep = 0; i_dep < number_of_reactants_[i_rxn]; ++i_dep)
{
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
v_jacobian[offset_jacobian + *flat_id + i_cell] += d_rate_d_ind[i_cell];
++flat_id;
}
for (std::size_t i_dep = 0; i_dep < number_of_products_[i_rxn]; ++i_dep)
{
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
v_jacobian[offset_jacobian + *flat_id + i_cell] -= yield[i_dep] * d_rate_d_ind[i_cell];
++flat_id;
}
const std::size_t idx_state_variables = offset_state + (react_id[i_react] * L);
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
d_rate_d_ind[i_cell] *= v_state_variables[idx_state_variables + i_cell];
}
react_id += number_of_reactants_[i_rxn];
yield += number_of_products_[i_rxn];
for (std::size_t i_dep = 0; i_dep < process_info.number_of_dependent_reactants_+1; ++i_dep)
{
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
v_jacobian[offset_jacobian + *flat_id + i_cell] += d_rate_d_ind[i_cell];
++flat_id;
}
for (std::size_t i_dep = 0; i_dep < process_info.number_of_products_; ++i_dep)
{
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
v_jacobian[offset_jacobian + *flat_id + i_cell] -= yield[i_dep] * d_rate_d_ind[i_cell];
++flat_id;
}
react_id += process_info.number_of_dependent_reactants_;
yield += process_info.number_of_products_;
}
}
}
Expand Down
66 changes: 39 additions & 27 deletions test/unit/process/test_process_set_policy.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,12 @@ void testProcessSet()

micm::Process r3 = micm::Process::Create().SetReactants({ quz }).SetProducts({}).SetPhase(gas_phase);

auto used_species = RatesPolicy::SpeciesUsed(std::vector<micm::Process>{ r1, r2, r3 });
micm::Process r4 = micm::Process::Create()
.SetReactants({ baz, qux })
.SetProducts({ Yields(bar, 1), Yields(quz, 2.5)})
.SetPhase(gas_phase);

auto used_species = RatesPolicy::SpeciesUsed(std::vector<micm::Process>{ r1, r2, r3, r4 });

EXPECT_EQ(used_species.size(), 6);
EXPECT_TRUE(used_species.contains("foo"));
Expand All @@ -55,37 +60,41 @@ void testProcessSet()
EXPECT_TRUE(used_species.contains("qux"));
EXPECT_FALSE(used_species.contains("corge"));

RatesPolicy set = RatesPolicy(std::vector<micm::Process>{ r1, r2, r3 }, state.variable_map_);
RatesPolicy set = RatesPolicy(std::vector<micm::Process>{ r1, r2, r3, r4 }, state.variable_map_);

EXPECT_EQ(state.variables_.NumRows(), 2);
EXPECT_EQ(state.variables_.NumColumns(), 6);
state.variables_[0] = { 0.1, 0.2, 0.3, 0.4, 0.5, 0.0 };
state.variables_[1] = { 1.1, 1.2, 1.3, 1.4, 1.5, 0.0 };
DenseMatrixPolicy rate_constants{ 2, 3 };
rate_constants[0] = { 10.0, 20.0, 30.0 };
rate_constants[1] = { 110.0, 120.0, 130.0 };
state.conditions_[0].air_density_ = 70.0;
state.conditions_[1].air_density_ = 80.0;
DenseMatrixPolicy rate_constants{ 2, 4 };
// the rate constants will have been calculated and combined with
// parameterized species before calculating forcing terms
rate_constants[0] = { 10.0, 20.0 * 70.0 * 0.72, 30.0, 40.0 * 70.0 * 0.72};
rate_constants[1] = { 110.0, 120.0 * 80.0 * 0.72, 130.0, 140.0 * 80.0 * 0.72};

DenseMatrixPolicy forcing{ 2, 5, 1000.0 };

set.template AddForcingTerms<DenseMatrixPolicy>(rate_constants, state.variables_, forcing);
EXPECT_EQ(forcing[0][0], 1000.0 - 10.0 * 0.1 * 0.3 + 20.0 * 0.2);
EXPECT_EQ(forcing[1][0], 1000.0 - 110.0 * 1.1 * 1.3 + 120.0 * 1.2);
EXPECT_EQ(forcing[0][1], 1000.0 + 10.0 * 0.1 * 0.3 - 20.0 * 0.2);
EXPECT_EQ(forcing[1][1], 1000.0 + 110.0 * 1.1 * 1.3 - 120.0 * 1.2);
EXPECT_EQ(forcing[0][2], 1000.0 - 10.0 * 0.1 * 0.3);
EXPECT_EQ(forcing[1][2], 1000.0 - 110.0 * 1.1 * 1.3);
EXPECT_EQ(forcing[0][3], 1000.0 + 20.0 * 0.2 * 1.4 - 30.0 * 0.4);
EXPECT_EQ(forcing[1][3], 1000.0 + 120.0 * 1.2 * 1.4 - 130.0 * 1.4);
EXPECT_EQ(forcing[0][4], 1000.0 + 10.0 * 0.1 * 0.3 * 2.4);
EXPECT_EQ(forcing[1][4], 1000.0 + 110.0 * 1.1 * 1.3 * 2.4);
EXPECT_DOUBLE_EQ(forcing[0][0], 1000.0 - 10.0 * 0.1 * 0.3 + 20.0 * 70.0 * 0.72 * 0.2); // foo
EXPECT_DOUBLE_EQ(forcing[1][0], 1000.0 - 110.0 * 1.1 * 1.3 + 120.0 * 80.0 * 0.72 * 1.2);
EXPECT_DOUBLE_EQ(forcing[0][1], 1000.0 + 10.0 * 0.1 * 0.3 - 20.0 * 0.2 * 70.0 * 0.72 + 40.0 * 70.0 * 0.72 * 0.3); // bar
EXPECT_DOUBLE_EQ(forcing[1][1], 1000.0 + 110.0 * 1.1 * 1.3 - 120.0 * 1.2 * 80.0 * 0.72 + 140.0 * 80.0 * 0.72 * 1.3);
EXPECT_DOUBLE_EQ(forcing[0][2], 1000.0 - 10.0 * 0.1 * 0.3 - 40.0 * 70.0 * 0.72 * 0.3); // baz
EXPECT_DOUBLE_EQ(forcing[1][2], 1000.0 - 110.0 * 1.1 * 1.3 - 140.0 * 80.0 * 0.72 * 1.3);
EXPECT_DOUBLE_EQ(forcing[0][3], 1000.0 + 20.0 * 70.0 * 0.72 * 0.2 * 1.4 - 30.0 * 0.4 + 40.0 * 70.0 * 0.72 * 2.5 * 0.3); // quz
EXPECT_DOUBLE_EQ(forcing[1][3], 1000.0 + 120.0 * 80.0 * 0.72 * 1.2 * 1.4 - 130.0 * 1.4 + 140.0 * 80.0 * 0.72 * 2.5 * 1.3);
EXPECT_DOUBLE_EQ(forcing[0][4], 1000.0 + 10.0 * 0.1 * 0.3 * 2.4); // quuz
EXPECT_DOUBLE_EQ(forcing[1][4], 1000.0 + 110.0 * 1.1 * 1.3 * 2.4);

auto non_zero_elements = set.NonZeroJacobianElements();
// ---- foo bar baz quz quuz
// foo 0 1 2 - -
// bar 3 4 5 - -
// baz 6 - 7 - -
// quz - 8 - 9 -
// quuz 10 - 11 - -
// quz - 8 9 10 -
// quuz 11 - 12 - -

auto elem = non_zero_elements.begin();
compare_pair(*elem, index_pair(0, 0));
Expand All @@ -97,6 +106,7 @@ void testProcessSet()
compare_pair(*(++elem), index_pair(2, 0));
compare_pair(*(++elem), index_pair(2, 2));
compare_pair(*(++elem), index_pair(3, 1));
compare_pair(*(++elem), index_pair(3, 2));
compare_pair(*(++elem), index_pair(3, 3));
compare_pair(*(++elem), index_pair(4, 0));
compare_pair(*(++elem), index_pair(4, 2));
Expand All @@ -109,22 +119,24 @@ void testProcessSet()
set.SubtractJacobianTerms(rate_constants, state.variables_, jacobian);
EXPECT_DOUBLE_EQ(jacobian[0][0][0], 100.0 + 10.0 * 0.3); // foo -> foo
EXPECT_DOUBLE_EQ(jacobian[1][0][0], 100.0 + 110.0 * 1.3);
EXPECT_DOUBLE_EQ(jacobian[0][0][1], 100.0 - 20.0); // foo -> bar
EXPECT_DOUBLE_EQ(jacobian[1][0][1], 100.0 - 120.0);
EXPECT_DOUBLE_EQ(jacobian[0][0][1], 100.0 - 20.0 * 70.0 * 0.72); // foo -> bar
EXPECT_DOUBLE_EQ(jacobian[1][0][1], 100.0 - 120.0 * 80.0 * 0.72);
EXPECT_DOUBLE_EQ(jacobian[0][0][2], 100.0 + 10.0 * 0.1); // foo -> baz
EXPECT_DOUBLE_EQ(jacobian[1][0][2], 100.0 + 110.0 * 1.1);
EXPECT_DOUBLE_EQ(jacobian[0][1][0], 100.0 - 10.0 * 0.3); // bar -> foo
EXPECT_DOUBLE_EQ(jacobian[1][1][0], 100.0 - 110.0 * 1.3);
EXPECT_DOUBLE_EQ(jacobian[0][1][1], 100.0 + 20.0); // bar -> bar
EXPECT_DOUBLE_EQ(jacobian[1][1][1], 100.0 + 120.0);
EXPECT_DOUBLE_EQ(jacobian[0][1][2], 100.0 - 10.0 * 0.1); // bar -> baz
EXPECT_DOUBLE_EQ(jacobian[1][1][2], 100.0 - 110.0 * 1.1);
EXPECT_DOUBLE_EQ(jacobian[0][1][1], 100.0 + 20.0 * 70.0 * 0.72); // bar -> bar
EXPECT_DOUBLE_EQ(jacobian[1][1][1], 100.0 + 120.0 * 80.0 * 0.72);
EXPECT_DOUBLE_EQ(jacobian[0][1][2], 100.0 - 10.0 * 0.1 - 40.0 * 70.0 * 0.72); // bar -> baz
EXPECT_DOUBLE_EQ(jacobian[1][1][2], 100.0 - 110.0 * 1.1 - 140.0 * 80.0 * 0.72);
EXPECT_DOUBLE_EQ(jacobian[0][2][0], 100.0 + 10.0 * 0.3); // baz -> foo
EXPECT_DOUBLE_EQ(jacobian[1][2][0], 100.0 + 110.0 * 1.3);
EXPECT_DOUBLE_EQ(jacobian[0][2][2], 100.0 + 10.0 * 0.1); // baz -> baz
EXPECT_DOUBLE_EQ(jacobian[1][2][2], 100.0 + 110.0 * 1.1);
EXPECT_DOUBLE_EQ(jacobian[0][3][1], 100.0 - 1.4 * 20.0); // quz -> bar
EXPECT_DOUBLE_EQ(jacobian[1][3][1], 100.0 - 1.4 * 120.0);
EXPECT_DOUBLE_EQ(jacobian[0][2][2], 100.0 + 10.0 * 0.1 + 40.0 * 70.0 * 0.72); // baz -> baz
EXPECT_DOUBLE_EQ(jacobian[1][2][2], 100.0 + 110.0 * 1.1 + 140.0 * 80.0 * 0.72);
EXPECT_DOUBLE_EQ(jacobian[0][3][1], 100.0 - 1.4 * 20.0 * 70.0 * 0.72); // quz -> bar
EXPECT_DOUBLE_EQ(jacobian[1][3][1], 100.0 - 1.4 * 120.0 * 80.0 * 0.72);
EXPECT_DOUBLE_EQ(jacobian[0][3][2], 100.0 - 2.5 * 40.0 * 70.0 * 0.72); // quz -> baz
EXPECT_DOUBLE_EQ(jacobian[1][3][2], 100.0 - 2.5 * 140.0 * 80.0 * 0.72);
EXPECT_DOUBLE_EQ(jacobian[0][3][3], 100.0 + 30.0); // quz -> quz
EXPECT_DOUBLE_EQ(jacobian[1][3][3], 100.0 + 130.0);
EXPECT_DOUBLE_EQ(jacobian[0][4][0], 100.0 - 2.4 * 10.0 * 0.3); // quuz -> foo
Expand Down
Loading