void SetNonzeros<Add>::evalAdj(const std::vector<std::vector<MX> >& aseed,
std::vector<std::vector<MX> >& asens) {
// Get all the nonzeros
vector<int> nz = all();
// Number of derivative directions
int nadj = aseed.size();
// Output sparsity
const Sparsity &osp = sparsity();
const int* orow = osp.row();
vector<int> ocol = osp.get_col();
// Input sparsity (first input same as output)
const Sparsity &isp = dep(1).sparsity();
const int* irow = isp.row();
vector<int> icol = isp.get_col();
// We next need to resort the assignment vector by outputs instead of inputs
// Start by counting the number of output nonzeros corresponding to each input nonzero
vector<int> onz_count(osp.nnz()+2, 0);
for (vector<int>::const_iterator it=nz.begin(); it!=nz.end(); ++it) {
onz_count[*it+2]++;
}
// Cumsum to get index offset for output nonzero
for (int i=0; i<onz_count.size()-1; ++i) {
onz_count[i+1] += onz_count[i];
}
// Get the order of assignments
vector<int> nz_order(nz.size());
for (int k=0; k<nz.size(); ++k) {
// Save the new index
nz_order[onz_count[1+nz[k]]++] = k;
}
// Find out which elements are being set
vector<int>& with_duplicates = onz_count; // Reuse memory
onz_count.resize(nz.size());
for (int k=0; k<nz.size(); ++k) {
// Get output nonzero
int onz_k = nz[nz_order[k]];
// Get element (note: may contain duplicates)
if (onz_k>=0) {
with_duplicates[k] = ocol[onz_k]*osp.size1() + orow[onz_k];
} else {
with_duplicates[k] = -1;
}
}
// Get all output elements (this time without duplicates)
vector<int> el_output;
osp.find(el_output);
// Sparsity pattern being formed and corresponding nonzero mapping
vector<int> r_colind, r_row, r_nz, r_ind;
for (int d=0; d<nadj; ++d) {
// Get the matching nonzeros
r_ind.resize(el_output.size());
copy(el_output.begin(), el_output.end(), r_ind.begin());
aseed[d][0].sparsity().get_nz(r_ind);
// Sparsity pattern for the result
r_colind.resize(isp.size2()+1); // Col count
fill(r_colind.begin(), r_colind.end(), 0);
r_row.clear();
// Perform the assignments
r_nz.clear();
for (int k=0; k<nz.size(); ++k) {
// Get the corresponding nonzero for the input
int el = nz[k];
// Skip if zero assignment
if (el==-1) continue;
// Get the corresponding nonzero in the argument
int el_arg = r_ind[el];
// Skip if no argument
if (el_arg==-1) continue;
// Save the assignment
r_nz.push_back(el_arg);
// Get the corresponding element
int i=icol[k], j=irow[k];
// Add to sparsity pattern
r_row.push_back(j);
r_colind[1+i]++;
}
// col count -> col offset
for (int i=1; i<r_colind.size(); ++i) r_colind[i] += r_colind[i-1];
// If anything to set/add
if (!r_nz.empty()) {
// Create a sparsity pattern from vectors
Sparsity f_sp(isp.size1(), isp.size2(), r_colind, r_row);
asens[d][1] += aseed[d][0]->getGetNonzeros(f_sp, r_nz);
if (!Add) {
asens[d][0] += MX::zeros(f_sp)->getSetNonzeros(aseed[d][0], r_nz);
} else {
asens[d][0] += aseed[d][0];
}
} else {
asens[d][0] += aseed[d][0];
}
}
}
void GetNonzeros::evaluateMX(const MXPtrV& input, MXPtrV& output, const MXPtrVV& fwdSeed,
MXPtrVV& fwdSens, const MXPtrVV& adjSeed, MXPtrVV& adjSens,
bool output_given) {
// Get all the nonzeros
vector<int> nz = getAll();
// Number of derivative directions
int nfwd = fwdSens.size();
int nadj = adjSeed.size();
// Output sparsity
const Sparsity& osp = sparsity();
const vector<int>& orow = osp.row();
vector<int> ocol = osp.getCol();
// Input sparsity
const Sparsity& isp = dep().sparsity();
//const vector<int>& irow = isp.row();
vector<int> icol = isp.getCol();
// Get all input elements
vector<int> el_input;
isp.getElements(el_input, false);
// Sparsity pattern being formed and corresponding nonzero mapping
vector<int> r_colind, r_row, r_nz, r_ind;
// Nondifferentiated function and forward sensitivities
int first_d = output_given ? 0 : -1;
for (int d=first_d; d<nfwd; ++d) {
// Get references to arguments and results
const MX& arg = d<0 ? *input[0] : *fwdSeed[d][0];
MX& res = d<0 ? *output[0] : *fwdSens[d][0];
// Get the matching nonzeros
r_ind.resize(el_input.size());
copy(el_input.begin(), el_input.end(), r_ind.begin());
arg.sparsity().getNZInplace(r_ind);
// Sparsity pattern for the result
r_colind.resize(osp.size2()+1); // Col count
fill(r_colind.begin(), r_colind.end(), 0);
r_row.clear();
// Perform the assignments
r_nz.clear();
for (int k=0; k<nz.size(); ++k) {
// Get the corresponding nonzero for the input
int el = nz[k];
// Skip if zero assignment
if (el==-1) continue;
// Get the corresponding nonzero in the argument
int el_arg = r_ind[el];
// Skip if no argument
if (el_arg==-1) continue;
// Save the assignment
r_nz.push_back(el_arg);
// Get the corresponding element
int i=ocol[k], j=orow[k];
// Add to sparsity pattern
r_row.push_back(j);
r_colind[1+i]++;
}
// col count -> col offset
for (int i=1; i<r_colind.size(); ++i) r_colind[i] += r_colind[i-1];
// Create a sparsity pattern from vectors
if (r_nz.size()==0) {
res = MX::sparse(osp.shape());
} else {
Sparsity f_sp(osp.size1(), osp.size2(), r_colind, r_row);
res = arg->getGetNonzeros(f_sp, r_nz);
}
}
// Adjoint sensitivities
for (int d=0; d<nadj; ++d) {
// Get an owning references to the seeds and sensitivities
// and clear the seeds for the next run
MX aseed = *adjSeed[d][0];
*adjSeed[d][0] = MX();
MX& asens = *adjSens[d][0]; // Sensitivity after addition
MX asens0 = asens; // Sensitivity before addition
// Get the corresponding nz locations in the output sparsity pattern
aseed.sparsity().getElements(r_nz, false);
osp.getNZInplace(r_nz);
// Filter out ignored entries and check if there is anything to add at all
bool elements_to_add = false;
for (vector<int>::iterator k=r_nz.begin(); k!=r_nz.end(); ++k) {
if (*k>=0) {
if (nz[*k]>=0) {
elements_to_add = true;
} else {
*k = -1;
}
}
}
// Quick continue of no elements to add
if (!elements_to_add) continue;
// Get the nz locations in the adjoint sensitivity corresponding to the inputs
r_ind.resize(el_input.size());
copy(el_input.begin(), el_input.end(), r_ind.begin());
asens0.sparsity().getNZInplace(r_ind);
// Enlarge the sparsity pattern of the sensitivity if not all additions fit
for (vector<int>::iterator k=r_nz.begin(); k!=r_nz.end(); ++k) {
if (*k>=0 && r_ind[nz[*k]]<0) {
// Create a new pattern which includes both the the previous seed and the addition
Sparsity sp = asens0.sparsity().patternUnion(dep().sparsity());
asens0 = asens0->getSetSparse(sp);
// Recalculate the nz locations in the adjoint sensitivity corresponding to the inputs
copy(el_input.begin(), el_input.end(), r_ind.begin());
asens0.sparsity().getNZInplace(r_ind);
break;
}
}
// Have r_nz point to locations in the sensitivity instead of the output
for (vector<int>::iterator k=r_nz.begin(); k!=r_nz.end(); ++k) {
if (*k>=0) {
*k = r_ind[nz[*k]];
}
}
// Add to the element to the sensitivity
asens = aseed->getAddNonzeros(asens0, r_nz);
}
}
