Commit 1855da17 authored by Jan Oliver Oelerich's avatar Jan Oliver Oelerich

Added center of mass calculation

parent 778e9f37
......@@ -83,6 +83,26 @@ namespace stemsalabim {
_adf_intensities_stored = true;
}
/*!
* Stores COM in the member vector.
* @param in COM as a vector of float
* @param buf memory buffer where the data is stored
*/
void storeCOM(const std::vector<float> &in, ptr_buf_t &buf) {
com = buf->add(in);
_com_stored = true;
}
/*!
* Stores COM in the member vector given a pre-created buffer entry.
* @param in COM as a vector of float
* @param buf memory buffer where the data is stored
*/
void storeCOM(memory::buffer::entry<float> & buf_entry) {
com = buf_entry;
_com_stored = true;
}
/*!
* Stores CBED intensities in the member vector.
* @param in CBED intensities as a vector of float
......@@ -113,6 +133,15 @@ namespace stemsalabim {
_cbed_intensities_stored = true;
}
/*!
* Clear COM in the member vector.
* @param buf memory buffer where the data is stored
*/
void clearCOM(ptr_buf_t &buf) {
buf->remove(com);
_com_stored = false;
}
/*!
* Clear intensities in the member vector.
* @param buf memory buffer where the data is stored
......@@ -147,6 +176,14 @@ namespace stemsalabim {
return _adf_intensities_stored;
}
/*!
* Check if COM data are stored in this struct.
* @return true if COM are stored.
*/
bool hasCOM() const {
return _com_stored;
}
// members are public, so we don't need getters/setters.
unsigned int adf_row;
......@@ -159,9 +196,11 @@ namespace stemsalabim {
bool adf{false};
bool cbed{false};
bool _com_stored{false};
bool _adf_intensities_stored{false};
bool _cbed_intensities_stored{false};
memory::buffer::entry<float> com;
memory::buffer::entry<float> adf_intensities;
memory::buffer::entry<float> cbed_intensities;
......
......@@ -32,7 +32,7 @@
using namespace std;
using namespace stemsalabim;
void IO::initNcFile(const shared_ptr<GridManager> &gridman, const shared_ptr<atomic::Cell> &crystal) {
void IO::initNCFile(const shared_ptr<GridManager> &gridman, const shared_ptr<atomic::Cell> &crystal) {
Params &p = Params::getInstance();
auto &mpi_env = mpi::Environment::getInstance();
......@@ -167,8 +167,10 @@ void IO::initNcFile(const shared_ptr<GridManager> &gridman, const shared_ptr<ato
auto defocusDim = g.defineDim("adf_defocus", p.adfAverageDefoci() ? 1 : p.numberOfDefoci());
auto fpDim = g.defineDim("adf_phonon", p.adfAverageConfigurations() ? 1 : p.numberOfConfigurations());
auto sliceDim = g.defineDim("adf_slice", gridman->adfSliceCoords().size());
auto coordinateDim = g.defineDim("coordinate_dim", 2);
auto v = g.defineVar<float>("adf_intensities", vd({defocusDim, xDim, yDim, fpDim, sliceDim, angleDim}));
g.defineVar<float>("center_of_mass", vd({defocusDim, xDim, yDim, fpDim, sliceDim, coordinateDim}));
g.defineVar<double>("adf_probe_x_grid", vd({xDim}));
g.defineVar<double>("adf_probe_y_grid", vd({yDim}));
g.defineVar<double>("adf_detector_grid", vd({angleDim}));
......@@ -226,6 +228,7 @@ void IO::initBuffers(const shared_ptr<GridManager> &gridman) {
if(p.adf()) {
_adf_intensities_buffer.resize(gridman->adfDetectorGrid().size() * gridman->adfSliceCoords().size());
_com_buffer.resize(2 * gridman->adfSliceCoords().size());
}
if(p.cbed()) {
......@@ -444,6 +447,85 @@ void IO::writeCrystal(shared_ptr<const FPConfManager> fpman, shared_ptr<const at
}
}
void
IO::writeCOM(unsigned int idefocus, unsigned int iconf, const shared_ptr<GridManager> &gridman, const ScanPoint &point,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf) {
unique_lock<mutex> qlck(_io_lock);
Params &p = Params::getInstance();
if(!point.adf || !point.hasCOM())
return;
NCFile f(p.outputFilename());
auto g = f.group("adf");
auto v = g.var("center_of_mass");
try {
IO::writeCOM(idefocus, iconf, gridman, point, ibuf, v);
} catch(NcException &e) {
output::error("STEM %s\n", e.what());
}
}
void
IO::writeCOM(unsigned int idefocus, unsigned int iconf, const shared_ptr<GridManager> &gridman, const ScanPoint &point,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf, NCVar &v) {
Params &p = Params::getInstance();
try {
unsigned int n_slices_total = (unsigned int) gridman->adfSliceCoords().size();
// intensities are ordered by rank and therefore internally by x and y and rad.
// we can write all data for one frozen phonon configuration at once.
if(p.adfAverageDefoci() && (iconf > 0 || idefocus > 0)) {
// see http://stackoverflow.com/questions/9915653/ for how to calculate a weighted, running average
double w = gridman->defocusWeights()[idefocus];
v.get(vs({0, point.adf_row, point.adf_col, 0, 0, 0}), vs({1, 1, 1, 1, n_slices_total, 2}), _com_buffer);
for(size_t j = 0; j < _com_buffer.size(); ++j)
_com_buffer[j] = (_com_buffer[j] * _running_weight + w * ibuf->value(point.com, j)) /
(_running_weight + w);
v.put(vs({0, point.adf_row, point.adf_col, 0, 0, 0}), vs({1, 1, 1, 1, n_slices_total, 2}), _com_buffer);
_running_weight += w;
} else if(!p.adfAverageDefoci() && p.adfAverageConfigurations() && iconf > 0) {
// no defocus average but configuration average
v.get(vs({idefocus, point.adf_row, point.adf_col, 0, 0, 0}),
vs({1, 1, 1, 1, n_slices_total, 2}),
_com_buffer);
for(size_t j = 0; j < _com_buffer.size(); ++j)
_com_buffer[j] = (_com_buffer[j] * iconf + ibuf->value(point.com, j)) / (iconf + 1);
v.put(vs({idefocus, point.adf_row, point.adf_col, 0, 0, 0}),
vs({1, 1, 1, 1, n_slices_total, 2}),
_com_buffer);
} else {
// no average, either when this is the first configuration over all or no averaging is chosen.
v.put(vs({idefocus, point.adf_row, point.adf_col, iconf, 0, 0}),
vs({1, 1, 1, 1, n_slices_total, 2}),
ibuf->ptr(point.com));
}
} catch(NcException &e) {
output::error("STEM %s\n", e.what());
}
}
void IO::writeAdfIntensities(unsigned int idefocus, unsigned int iconf, const shared_ptr<GridManager> &gridman,
const ScanPoint &point, std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf) {
unique_lock<mutex> qlck(_io_lock);
......@@ -617,6 +699,7 @@ string getTempFileName() {
}
void IO::writeTemporaryResult(unsigned int idefocus, unsigned int iconf, ScanPoint &point,
shared_ptr<memory::buffer::number_buffer<float>> &combuf,
shared_ptr<memory::buffer::number_buffer<float>> &ibuf,
shared_ptr<memory::buffer::number_buffer<float>> &cbuf) {
......@@ -647,6 +730,9 @@ void IO::writeTemporaryResult(unsigned int idefocus, unsigned int iconf, ScanPoi
tmpfile.write(reinterpret_cast<const char *>(&point.index), sizeof(point.index));
if(point.hasAdfIntensities()) {
tmpfile.write(reinterpret_cast<const char *>(combuf->ptr(point.com)), point.com.byte_size());
point.clearCOM(combuf);
tmpfile.write(reinterpret_cast<const char *>(ibuf->ptr(point.adf_intensities)),
point.adf_intensities.byte_size());
point.clearAdfIntensities(ibuf);
......@@ -663,6 +749,7 @@ void IO::writeTemporaryResult(unsigned int idefocus, unsigned int iconf, ScanPoi
}
void IO::copyFromTemporaryFile(const shared_ptr<GridManager> &gridman,
shared_ptr<memory::buffer::number_buffer<float>> &combuf,
shared_ptr<memory::buffer::number_buffer<float>> &ibuf,
shared_ptr<memory::buffer::number_buffer<float>> &cbuf) {
......@@ -684,11 +771,12 @@ void IO::copyFromTemporaryFile(const shared_ptr<GridManager> &gridman,
NCFile f(prms.outputFilename());
NCVar av, cv;
NCVar av, cv, comv;
try {
if(prms.adf()) {
av = f.group("adf").var("adf_intensities");
comv = f.group("adf").var("center_of_mass");
}
if(prms.cbed()) {
cv = f.group("cbed").var("cbed_intensities");
......@@ -709,6 +797,12 @@ void IO::copyFromTemporaryFile(const shared_ptr<GridManager> &gridman,
ScanPoint p = gridman->scanPoints()[pix_idx];
if(prms.adf() && p.adf) {
auto &com_entry = combuf->add_empty();
tmpfile.read(reinterpret_cast<char *>(combuf->ptr(com_entry)), com_entry.byte_size());
p.storeCOM(com_entry);
writeCOM(idefocus, iconf, gridman, p, combuf, comv);
p.clearCOM(ibuf);
auto &entry = ibuf->add_empty();
tmpfile.read(reinterpret_cast<char *>(ibuf->ptr(entry)), entry.byte_size());
p.storeAdfIntensities(entry);
......
......@@ -62,7 +62,7 @@ namespace stemsalabim {
* is written to the file.
* @param gridman Reference to the GridManager required for writing some information.
*/
void initNcFile(const std::shared_ptr<GridManager> &gridman, const std::shared_ptr<atomic::Cell> &crystal);
void initNCFile(const std::shared_ptr<GridManager> &gridman, const std::shared_ptr<atomic::Cell> &crystal);
/*!
* Initialize the memory buffers that are used when averaging results.
......@@ -102,6 +102,21 @@ namespace stemsalabim {
const std::shared_ptr<GridManager> &gridman, const ScanPoint &point,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf, NCVar & v);
/*!
* Append to the Center of mass variable.
* @param iconf current FP configuration index
* @param idefocus current defocus index
* @param gridman reference to a GridManager
* @param point The pixel for which COM values are to be written.
* @param ibuf The number buffer with the actual COM values.
*/
void writeCOM(unsigned int idefocus, unsigned int iconf,
const std::shared_ptr<GridManager> &gridman, const ScanPoint &point,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf);
void writeCOM(unsigned int idefocus, unsigned int iconf,
const std::shared_ptr<GridManager> &gridman, const ScanPoint &point,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf, NCVar & v);
/*!
* Append to the CBED intensities variable.
* @param iconf current FP configuration index
......@@ -127,6 +142,7 @@ namespace stemsalabim {
*/
static void writeTemporaryResult(unsigned int idefocus, unsigned int iconf,
ScanPoint &point,
std::shared_ptr<memory::buffer::number_buffer<float>> &combuf,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf,
std::shared_ptr<memory::buffer::number_buffer<float>> &cbuf);
......@@ -137,6 +153,7 @@ namespace stemsalabim {
* @param cbuf The number buffer with the actual cbed intensity values.
*/
void copyFromTemporaryFile(const std::shared_ptr<GridManager> &gridman,
std::shared_ptr<memory::buffer::number_buffer<float>> &combuf,
std::shared_ptr<memory::buffer::number_buffer<float>> &ibuf,
std::shared_ptr<memory::buffer::number_buffer<float>> &cbuf);
......@@ -176,6 +193,7 @@ namespace stemsalabim {
double _running_weight{0};
std::vector<float> _com_buffer;
std::vector<float> _adf_intensities_buffer;
std::vector<float> _cbed_intensities_buffer;
......
......@@ -92,7 +92,7 @@ void Simulation::init() {
// The crystal instance will be initialized here and then broadcast to
// tge slaves.
string config_file_contents;
_io->initNcFile(_gridman, p.cell());
_io->initNCFile(_gridman, p.cell());
{
NCFile f(p.paramsFileName(), false, true);
config_file_contents = f.group("params").att("config_file_contents");
......@@ -188,7 +188,7 @@ void Simulation::run() {
for(auto &t : threads)
t.join();
} else {
for(auto & slic: _gridman->slices()) {
for(auto &slic: _gridman->slices()) {
slic->calculateTransmissionFunction(_gridman, _io->readStoredPotential(_fpman, slic->id()));
}
}
......@@ -213,7 +213,7 @@ void Simulation::run() {
for(int rank = 0; rank < mpi_env.size(); rank++) {
if(rank == mpi_env.rank()) {
auto start = high_resolution_clock::now();
_io->copyFromTemporaryFile(_gridman, _adf_intensity_buffer, _cbed_intensity_buffer);
_io->copyFromTemporaryFile(_gridman, _com_buffer, _adf_intensity_buffer, _cbed_intensity_buffer);
printLine(st,
output::fmt("Finished copying pixels for %s.\n",
......@@ -245,11 +245,13 @@ void Simulation::calculatePixel(ScanPoint &point, const double defocus) {
Params &p = Params::getInstance();
vector<float> adf_intensities;
vector<float> center_of_mass;
vector<float> cbed_intensities;
vector<float> detector_intensities;
if(point.adf) {
adf_intensities.reserve(_gridman->adfSliceCoords().size() * _gridman->adfDetectorGrid().size());
center_of_mass.reserve(_gridman->adfSliceCoords().size() * 2);
detector_intensities.resize(_gridman->adfDetectorGrid().size());
}
......@@ -302,6 +304,21 @@ void Simulation::calculatePixel(ScanPoint &point, const double defocus) {
}
adf_intensities.insert(adf_intensities.end(), detector_intensities.begin(), detector_intensities.end());
// calculate center of mass
double total_intensity = wave.integratedIntensity(), cx = 0, cy = 0;
for(unsigned int ix = 0; ix < _gridman->samplingX(); ix++) {
for(unsigned int iy = 0; iy < _gridman->samplingY(); iy++) {
auto intens = (float) pow(abs(wave(ix, iy)), 2);
cx += intens * _gridman->kx(ix);
cy += intens * _gridman->kx(iy);
}
}
cx /= total_intensity;
cy /= total_intensity;
center_of_mass.push_back((float) (cx * p.wavelength() * 1e3));
center_of_mass.push_back((float) (cy * p.wavelength() * 1e3));
}
if(point.cbed && _gridman->cbedStoreSlice(slice->id())) {
......@@ -360,8 +377,10 @@ void Simulation::calculatePixel(ScanPoint &point, const double defocus) {
wave.backwardFFT();
}
if(point.adf)
if(point.adf) {
point.storeAdfIntensities(adf_intensities, _adf_intensity_buffer);
point.storeCOM(center_of_mass, _com_buffer);
}
if(point.cbed)
point.storeCBEDIntensities(cbed_intensities, _cbed_intensity_buffer);
......@@ -475,6 +494,7 @@ void Simulation::multisliceMaster(const SimulationState &st) {
IO::writeTemporaryResult(st.idefocus(),
st.iconf(),
sp,
_com_buffer,
_adf_intensity_buffer,
_cbed_intensity_buffer);
}
......@@ -495,7 +515,12 @@ void Simulation::multisliceMaster(const SimulationState &st) {
for(unsigned int px_idx : scan_work.getDirtyAndClean()) {
ScanPoint &sp = work_packages[px_idx];
IO::writeTemporaryResult(st.idefocus(), st.iconf(), sp, _adf_intensity_buffer, _cbed_intensity_buffer);
IO::writeTemporaryResult(st.idefocus(),
st.iconf(),
sp,
_com_buffer,
_adf_intensity_buffer,
_cbed_intensity_buffer);
}
scan_work.waitUntilFinished();
......@@ -623,7 +648,12 @@ void Simulation::multisliceWorker(const SimulationState &st) {
for(unsigned int px_idx : package) {
ScanPoint &sp = work_packages[px_idx];
IO::writeTemporaryResult(st.idefocus(), st.iconf(), sp, _adf_intensity_buffer, _cbed_intensity_buffer);
IO::writeTemporaryResult(st.idefocus(),
st.iconf(),
sp,
_com_buffer,
_adf_intensity_buffer,
_cbed_intensity_buffer);
}
}
......@@ -646,6 +676,7 @@ void Simulation::initBuffers() {
_gridman->storedCbedSizeY() *
_gridman->cbedSliceCoords().size()) : 0;
_com_buffer = std::make_shared<buf_type>(2 * _gridman->adfSliceCoords().size(), 2 * prms.workPackageSize());
_adf_intensity_buffer = std::make_shared<buf_type>(number_intensities_per_pixel, 2 * prms.workPackageSize());
_cbed_intensity_buffer = std::make_shared<buf_type>(number_cbed_per_pixel, 2 * prms.workPackageSize());
}
......@@ -659,6 +690,9 @@ void Simulation::storeFinishedPixels(const SimulationState &st, vector<ScanPoint
ScanPoint &point = work_packages[index];
if(prms.adf() && point.adf && point.hasAdfIntensities()) {
_io->writeCOM(st.idefocus(), st.iconf(), _gridman, point, _com_buffer);
point.clearCOM(_com_buffer);
_io->writeAdfIntensities(st.idefocus(), st.iconf(), _gridman, point, _adf_intensity_buffer);
point.clearAdfIntensities(_adf_intensity_buffer);
}
......
......@@ -101,6 +101,7 @@ namespace stemsalabim {
/// these buffers temporarily hold intensity data as continuous storage. This is to prevent
/// memory fragmentation errors.
std::shared_ptr<memory::buffer::number_buffer<float>> _adf_intensity_buffer;
std::shared_ptr<memory::buffer::number_buffer<float>> _com_buffer;
std::shared_ptr<memory::buffer::number_buffer<float>> _cbed_intensity_buffer;
/*!
......
......@@ -21,7 +21,6 @@
#include <iostream>
#include "utilities/Wave.hpp"
#include "classes/Crystal.hpp"
#include "classes/Simulation.hpp"
using namespace std;
......
......@@ -133,7 +133,7 @@ int main(int argc, const char **argv) {
fpman->generateDisplacements();
gridman->generateGrids();
io->initNcFile(gridman, p.cell());
io->initNCFile(gridman, p.cell());
io->writeParams(gridman);
io->writeGrids(gridman);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment