mp2prep.C

Go to the documentation of this file.

/*
    BM@N alignment routine
    BM@N experiment at NICA complex, JINR, 2025
 
    Department: Math & Soft Group of HEP lab
    Author:     Igor Polev, polev@jinr.int
 
    Millepede-II input generator (cross-check tool for BmnAligner)
 
    Translates one BmnAligner task - measurement model, detector model,
    constraints, presigmas, dataset - into the four files pede consumes:
 
        mp2mille.bin   binary records: one local fit (track) per record
        mp2par.txt     initial values and pre-sigmas of the global parameters
        mp2con.txt     linear equality constraints
        mp2steer.txt   steering file referencing the three above
        mp2labels.txt  label -> object/parameter map (for res2json.py only)
 
    The task is hard-wired to BmnMeasureLin4 + BmnRun8TopDM (see run.sh); the
    config file paths are arguments.
*/
 
#include "BmnAlignDefines.h"
#include "BmnMeasureLin4.h"
#include "BmnRootIterator.h"
#include "BmnRun8TopDM.h"
#include "BmnWeightedLine.h"
#include "Mille.cc"   // millepede-ii ships the C binary writer as source only
#include "Mille.h"
#include "fairlogger/Logger.h"
 
#include <array>
#include <cmath>
#include <fstream>
#include <map>
#include <nlohmann/json.hpp>
#include <string>
#include <vector>
 
/*
 *  Parametrization
 *  ---------------
 *  BmnMeasureLin4 leaves phi and theta without effect, so - unlike BmnAligner,
 *  whose parameter vector is fixed at 6 and which has to fix the two by
 *  constraints - they simply do not exist here: pede only knows the parameters
 *  that appear in the binary file. Anything addressing them in the config files
 *  is dropped, with a report.
 *
 *  The two levels of BmnAlignHierarchy (38 elements, 11 station halves) are
 *  expressed the way Millepede does hierarchies: every measurement carries the
 *  derivative twice, once against its element's own parameter and once against
 *  the same parameter of the owning composite. That is exact for the four
 *  parameters kept here - the propagation Jacobian J of BmnAlignHierarchy is
 *  the identity for dX/dY/dZ/psi once the composite pivot sits on the beam
 *  axis, which is the pivot default the models use.
 */
namespace
{
 
constexpr Int_t kParams{4};
// index of each kept parameter within the 6-slot SVectGL of the measure model
constexpr std::array<Int_t, kParams> kSlot{idX, idY, idZ, iPsi};
constexpr std::array<const char*, kParams> kName{"dX", "dY", "dZ", "psi"};
 
// label = object number (1-based) * 10 + parameter slot (1-based), composites
// shifted out of the elements' range so that a label stays readable by eye
constexpr Int_t kStride{10}, kCompBase{100000};
 
inline Int_t LeafLabel(Int_t leaf, Int_t p)
{
    return (leaf + 1) * kStride + p + 1;
}
inline Int_t CompLabel(Int_t comp, Int_t p)
{
    return kCompBase + ((comp + 1) * kStride + p + 1);
}
 
// -1 for phi and theta: not parameters of this task
inline Int_t SlotByName(const std::string& name)
{
    for (Int_t p = 0; p < kParams; p++)
        if (name == kName[p])
            return p;
    return -1;
}
 
struct HitRec
{
    Double_t x, y, z, wx, wy;
    Int_t leaf;
};
 
// weighted straight-line fit, x(z) = alpha + beta * z, as in
// BmnAligner::LinRegressionAlphaPart but on the nominal (unaligned) geometry
Bool_t FitLine(const std::vector<HitRec>& hits, Bool_t inY, Double_t& alpha, Double_t& beta)
{
    BmnWeightedLine line;
    for (const HitRec& h : hits)
        line.Add(h.z, inY ? h.y : h.x, inY ? h.wy : h.wx);
    constexpr Double_t SINGULAR_DET{1.0E-6};   // hits too close in z to fit a line through
    if (std::abs(line.Det()) < SINGULAR_DET)
        return kFALSE;
    beta = line.Slope();
    alpha = line.At0();
    return kTRUE;
}
 
}   // namespace
 
int mp2prep(const char* fileListPath = "/home/igor/bmnroot/align_new/cfg/run8/7466.noa_rlerr.list",
            const char* constraintsPath = "/home/igor/bmnroot/align_new/cfg/run8/constraints/cnstr_38e_lin4.json",
            const char* presigmaPath = "/home/igor/bmnroot/align_new/cfg/run8/presigmas/psg_38e_lin4.json",
            const char* outDir = "/home/igor/bmnroot/align_new/mp2/out",
            Int_t maxTracks = 0)   // 0 means no limit
{
    fair::Logger::SetConsoleColor(true);
    fair::Logger::SetConsoleSeverity("info");
 
    BmnMeasureLin4 measure;
    BmnRun8TopDM model;
    if (!model.Init())
        return -1;
    LOG(info) << "    detector model  : " << model.Description();
    LOG(info) << "    measure model   : " << measure.Description();
 
    // --- hierarchy: which composite owns which element ---------------------
    const std::vector<BmnDetectorModel::Composite> composites{model.Composites()};
    VecInt_t compOf(model.Alignables(), -1);
    std::map<std::string, Int_t> compByName;
    for (std::size_t c = 0; c < composites.size(); c++) {
        compByName[composites[c].name] = c;
        for (Int_t leaf : composites[c].leaves)
            compOf[leaf] = c;
    }
    for (Int_t leaf = 0; leaf < model.Alignables(); leaf++)
        if (compOf[leaf] < 0) {
            LOG(error) << " <!> element " << leaf << " belongs to no composite";
            return -1;
        }
    LOG(info) << "    parametrization : " << model.Alignables() << " elements + " << composites.size()
              << " composites, " << kParams
              << " parameters each = " << (model.Alignables() + (Int_t)composites.size()) * kParams
              << " global parameters";
 
    const std::string out{outDir};
 
    // --- labels map ---------------------------------------------------------
    {
        std::ofstream ofs(out + "/mp2labels.txt");
        if (!ofs.good()) {
            LOG(error) << " <!> cannot write into " << out << " (does it exist?)";
            return -1;
        }
        ofs << "# label object param owner\n";
        for (Int_t leaf = 0; leaf < model.Alignables(); leaf++)
            for (Int_t p = 0; p < kParams; p++)
                ofs << LeafLabel(leaf, p) << ' ' << model.EncodedID(leaf) << ' ' << kName[p] << ' '
                    << composites[compOf[leaf]].name << '\n';
        for (std::size_t c = 0; c < composites.size(); c++)
            for (Int_t p = 0; p < kParams; p++)
                ofs << CompLabel(c, p) << ' ' << composites[c].name << ' ' << kName[p] << " -\n";
    }
 
    // --- presigmas -> parameter file ----------------------------------------
    {
        std::ifstream ifs(presigmaPath);
        if (!ifs.good()) {
            LOG(error) << " <!> failed to open presigmas file " << presigmaPath;
            return -1;
        }
        nlohmann::json js;
        ifs >> js;
 
        std::array<Double_t, kParams> leafSigma{};
        std::vector<std::array<Double_t, kParams>> compSigma(composites.size());
        Int_t dropped{0};
        for (const nlohmann::json& e : js) {
            if (!e.contains("ParamName") || !e.contains("Sigma")) {
                LOG(error) << " <!> presigma entry without ParamName/Sigma:"
                              " only the (ParamName[, NodeName], Sigma) form is supported";
                return -1;
            }
            const Int_t p{SlotByName(e["ParamName"].get<std::string>())};
            if (p < 0) {
                dropped++;
                continue;
            }   // phi/theta: not a parameter here
            const Double_t sigma{e["Sigma"].get<Double_t>()};
            if (!e.contains("NodeName")) {
                leafSigma[p] = sigma;   // blanket entry: every element
                continue;
            }
            const std::string node{e["NodeName"].get<std::string>()};
            const auto it{compByName.find(node)};
            if (it == compByName.end()) {
                LOG(error) << " <!> presigma NodeName '" << node << "' is not a composite of the model";
                return -1;
            }
            compSigma[it->second][p] = sigma;
        }
        LOG_IF(warning, dropped) << "  ! " << dropped << " presigma entrie(s) on phi/theta ignored";
 
        std::ofstream ofs(out + "/mp2par.txt");
        ofs << "* label   initial value   pre-sigma (0 = free)\n"
            << "Parameter\n";
        for (Int_t leaf = 0; leaf < model.Alignables(); leaf++)
            for (Int_t p = 0; p < kParams; p++)
                ofs << LeafLabel(leaf, p) << "  0.0  " << leafSigma[p] << '\n';
        for (std::size_t c = 0; c < composites.size(); c++)
            for (Int_t p = 0; p < kParams; p++)
                ofs << CompLabel(c, p) << "  0.0  " << compSigma[c][p] << '\n';
        LOG(info) << "    " << out << "/mp2par.txt : pre-sigmas written";
    }
 
    // --- constraints --------------------------------------------------------
    {
        std::ifstream ifs(constraintsPath);
        if (!ifs.good()) {
            LOG(error) << " <!> failed to open constraints file " << constraintsPath;
            return -1;
        }
        nlohmann::json js;
        ifs >> js;
 
        std::ofstream ofs(out + "/mp2con.txt");
        Int_t written{0}, skipped{0};
        for (const nlohmann::json& jc : js) {
            std::map<Int_t, Double_t> coef;
            Double_t rhs{0.0};
            Bool_t alive{kTRUE};
            std::string comment{jc.value("Comment", std::string{})};
 
            for (const nlohmann::json& jco : jc["Coefficients"]) {
                if (jco.contains("RHS")) {
                    rhs = jco["RHS"].get<Double_t>();
                    continue;
                }
                if (!jco.contains("ParamName") || !jco.contains("Value")) {
                    LOG(error) << " <!> constraint coefficient without ParamName/Value:"
                                  " only the (ParamName, DetectorID|NodeName, Value)"
                                  " form is supported";
                    return -1;
                }
                const Int_t p{SlotByName(jco["ParamName"].get<std::string>())};
                if (p < 0) {   // constrains phi or theta: nothing to constrain
                    alive = kFALSE;
                    break;
                }
                const Double_t value{jco["Value"].get<Double_t>()};
                if (jco.contains("NodeName")) {
                    const std::string node{jco["NodeName"].get<std::string>()};
                    const auto it{compByName.find(node)};
                    if (it == compByName.end()) {
                        LOG(error) << " <!> constraint NodeName '" << node << "' is not a composite of the model";
                        return -1;
                    }
                    coef[CompLabel(it->second, p)] += value;
                } else if (jco.contains("DetectorID")) {
                    const Int_t id{jco["DetectorID"].get<Int_t>()};
                    if (model.UnknownID(id)) {
                        LOG(error) << " <!> constraint DetectorID " << id << " is not an element of the model";
                        return -1;
                    }
                    const Int_t leaf{model.Idx(id)};
                    // A coefficient on an element constrains its *effective*
                    // correction, so it must reach the composite too - this is
                    // the J^T c_leaf term of BmnAlignHierarchy::PromoteConstraints,
                    // with J the identity (see the note on the parametrization).
                    coef[LeafLabel(leaf, p)] += value;
                    coef[CompLabel(compOf[leaf], p)] += value;
                } else {
                    LOG(error) << " <!> constraint coefficient with neither NodeName"
                                  " nor DetectorID";
                    return -1;
                }
            }
            if (!alive || coef.empty()) {
                skipped++;
                LOG(info) << "      skipped (phi/theta only): " << comment;
                continue;
            }
            ofs << "\n* " << comment << "\nConstraint " << rhs << '\n';
            for (const auto& [label, value] : coef)
                ofs << "  " << label << "  " << value << '\n';
            written++;
        }
        LOG(info) << "    " << out << "/mp2con.txt : " << written << " constraint(s) written, " << skipped
                  << " skipped as meaningless for this model";
    }
 
    // --- steering -----------------------------------------------------------
    {
        std::ofstream ofs(out + "/mp2steer.txt");
        ofs << "* Millepede-II steering, generated by mp2prep.C\n"
            << "mp2par.txt\n"
            << "mp2con.txt\n"
            << "Cfiles\n"
            << "mp2mille.bin\n"
            << "\n"
            << "* full inversion: 196 parameters, and it is the covariance matrix\n"
            << "* that makes the result comparable with BmnAligner's reported errors\n"
            << "method inversion 5 0.001\n"
            << "end\n";
    }
 
    // --- binary records -----------------------------------------------------
    Mille mille((out + "/mp2mille.bin").c_str());
    BmnRootIterator itr(fileListPath, model, nullptr, kTRUE, maxTracks);
    if (!itr.Initialized()) {
        LOG(error) << " <!> failed to open the dataset " << fileListPath;
        return -1;
    }
 
    std::vector<HitRec> hits;
    Long64_t nTracks{0}, nHits{0}, nBadFit{0};
    itr.ResetAll();
    while (!itr.EndOfTracks()) {
        hits.clear();
        while (!itr.EndOfHits()) {
            const Int_t detID{itr.HitDetectorID()};
            if (model.UnknownID(detID)) {
                LOG(error) << " <!> unknown detector module encoding " << detID;
                return -1;
            }
            hits.push_back(HitRec{itr.HitX(), itr.HitY(), itr.HitZ(), itr.HitWx(), itr.HitWy(), model.Idx(detID)});
            itr.NextHit();
        }
 
        // Global derivatives depend on the track parameters, so the record needs
        // a reference trajectory. The measurements themselves go in unmodified:
        // the model is linear in the local parameters, so pede refits them.
        Double_t alphaX, betaX, alphaY, betaY;
        if (!FitLine(hits, kFALSE, alphaX, betaX) || !FitLine(hits, kTRUE, alphaY, betaY)) {
            nBadFit++;
            itr.NextTrack();
            continue;
        }
        const SVectLC lc{alphaX, betaX, alphaY, betaY};
 
        for (const HitRec& h : hits) {
            const SVectGL gx{measure.dXdA(h.z, lc)}, gy{measure.dYdA(h.z, lc)};
            const Int_t comp{compOf[h.leaf]};
 
            Int_t label[2 * kParams];
            float derGl[2 * kParams];
            for (Int_t p = 0; p < kParams; p++) {
                label[p] = LeafLabel(h.leaf, p);
                label[kParams + p] = CompLabel(comp, p);
            }
 
            float derLc[BMN_LOCAL_PARAMS_PT]{1.0f, (float)h.z, 0.0f, 0.0f};
            for (Int_t p = 0; p < kParams; p++)
                derGl[p] = derGl[kParams + p] = (float)gx(kSlot[p]);
            mille.mille(BMN_LOCAL_PARAMS_PT, derLc, 2 * kParams, derGl, label, (float)h.x,
                        (float)(1.0 / std::sqrt(h.wx)));
 
            derLc[0] = 0.0f;
            derLc[1] = 0.0f;
            derLc[2] = 1.0f;
            derLc[3] = (float)h.z;
            for (Int_t p = 0; p < kParams; p++)
                derGl[p] = derGl[kParams + p] = (float)gy(kSlot[p]);
            mille.mille(BMN_LOCAL_PARAMS_PT, derLc, 2 * kParams, derGl, label, (float)h.y,
                        (float)(1.0 / std::sqrt(h.wy)));
        }
        mille.end();
 
        nHits += hits.size();
        if (++nTracks % 500000 == 0)
            LOG(info) << "    " << nTracks << " tracks written...";
        itr.NextTrack();
    }
 
    LOG(info) << "    " << out << "/mp2mille.bin : " << nTracks << " records, " << nHits << " hits (" << 2 * nHits
              << " measurements)";
    LOG_IF(warning, nBadFit) << "  ! " << nBadFit << " track(s) skipped: degenerate line fit";
    LOG(state) << " < mp2prep : done, run pede from " << out;
    return 0;
}