Jungfraujoch/image_analysis/lattice_search/LatticeSearch.cpp

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "../../common/JFJochMath.h"
#include "LatticeSearch.h"
#include <gemmi/cellred.hpp>

#include <algorithm>
#include <cmath>
#include <optional>

// How close the reduced beta has to be to 90 degrees for the two Niggli types to be genuinely
// interchangeable (see the retry at the end of LatticeSearch). Not the angle tolerance: that is how
// far a metric may sit from an ideal one and still be called it, which is far too generous here - a
// cell 2 degrees off the boundary is a real type-1 cell, and presenting it in the obtuse setting
// promotes a general triclinic lattice to C-centred monoclinic on residuals of ~2 degrees. Measured:
// the crystal this was found on sits 0.07 degrees from the boundary and matches on 0.006 to 0.135;
// the triclinic cell that must not be promoted sits 2.0 degrees from it.
constexpr double NIGGLI_TYPE_BOUNDARY_DEG = 0.5;

struct NiggliClass {
    int number;
    int type;
    bool cond_AB;
    bool cond_BC;
    double cond_D;
    double cond_E;
    double cond_F;
    bool cond_DEF;
    bool cond_2DF;
    gemmi::Mat33 reindex;
    gemmi::CrystalSystem system;
    char centering;
};

namespace {
// The body of both entry points. only_class, when given, keeps just the characters of that Bravais
// class - see LatticeSearchForClass. With no filter this is the original walk unchanged, and the
// triclinic character fits every metric, so it always returns a result.
std::optional<LatticeSearchResult> SearchCharacters(const CrystalLattice &L, double dist_tolerance,
                                                    double angle_tolerance,
                                                    const std::pair<gemmi::CrystalSystem, char> *only_class) {
    UnitCell uc = L.GetUnitCell();
    gemmi::UnitCell g_uc(uc.a, uc.b, uc.c, uc.alpha, uc.beta, uc.gamma);
    // Niggli reduction following Krivy & Gruber (1976) Acta Cryst. A32, 297-298, via gemmi
    gemmi::GruberVector g_vec(g_uc, 'P', true);
    // The reduction decides the Niggli TYPE from the signs of the three scalar products, and gemmi's
    // default epsilon is 1e-9 ABSOLUTE while those products are 10^3 to 10^5 A^2 on a cell held in
    // float. A product that is structurally zero therefore arrives carrying ~1e-4 A^2 of rounding and
    // is read as definitely signed, the reduction lands on the wrong side of the type-I/type-II
    // boundary, and the character written for the other side matches nothing. Measured: a body-centred
    // tetragonal lattice with c > a*sqrt(2) loses its 4-fold on 38 of 60 rotations OF THE SAME LATTICE,
    // which is why an axis-aligned test never sees it and every cell the pipeline classifies is
    // refined and rotated.
    //
    // Scaling it by the cell's own magnitude puts the constant on a plateau three decades wide with
    // the over-call count flat across all of it. It is 100x the value Grosse-Kunstleve et al. give,
    // because theirs is calibrated for a double-precision cell and ours is float: measured, their
    // constant recovers 6% of these lattices and this one 93%.
    // Following Grosse-Kunstleve, Sauter & Adams (2004) Acta Cryst. A60, 1-6
    g_vec.niggli_reduce(1e-5 * std::max({g_vec.A, g_vec.B, g_vec.C}));

    CrystalLattice L_niggli = L;
    if (g_vec.change_of_basis)
        L_niggli = L.Multiply(gemmi::rot_as_mat33(g_vec.change_of_basis->rot).transpose());

    double A = g_vec.A;
    double B = g_vec.B;
    double C = g_vec.C;
    double D = g_vec.xi / 2;
    double E = g_vec.eta / 2;
    double F = g_vec.zeta / 2;

    // D, E, F are parameters so the table can also be built for the type-flipped setting below.
    // Lattice characters following International Tables for Crystallography Vol. A, Table 9.2.5.1
    auto make_classes = [&](double D, double E, double F) {
      return std::vector<NiggliClass>{
        {
            1, 1,
            true, true, A / 2, A / 2, A / 2, false, false,
            gemmi::Mat33{1, -1, 1, 1, 1, -1, -1, 1, 1},
            gemmi::CrystalSystem::Cubic, 'F'
        },
        {
            2, 1,
            true, true, D, D, D, false, false,
            {1, -1, 0, -1, 0, 1, -1, -1, -1},
            gemmi::CrystalSystem::Trigonal, 'R'
        },
        {
            3, 2,
            true, true, 0, 0, 0, false, false,
            gemmi::Mat33{1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Cubic, 'P'
        },
        {
            5, 2,
            true, true, -A / 3, -A / 3, -A / 3, false, false,
            gemmi::Mat33{1, 0, 1, 1, 1, 0, 0, 1, 1},
            gemmi::CrystalSystem::Cubic, 'I'
        },
        {
            4, 2,
            true, true, D, D, D, false, false,
            {1, -1, 0, -1, 0, 1, -1, -1, -1},
            gemmi::CrystalSystem::Trigonal, 'R'
        },
        {
            6, 2,
            true, true, D, D, F, true, false,
            {0, 1, 1, 1, 0, 1, 1, 1, 0},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            7, 2,
            true, true, D, E, E, true, false,
            {1, 0, 1, 1, 1, 0, 0, 1, 1},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            8, 2,
            true, true, D, E, F, true, false,
            {-1, -1, 0, -1, 0, -1, 0, -1, -1},
            gemmi::CrystalSystem::Orthorhombic, 'I'
        },
        {
            9, 1,
            true, false, A / 2, A / 2, A / 2, false, false,
            {1, 0, 0, -1, 1, 0, -1, -1, 3},
            gemmi::CrystalSystem::Trigonal, 'R'
        },

        {
            10, 1,
            true, false, D, D, F, false, false,
            {1, 1, 0, 1, -1, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            11, 2,
            true, false, 0, 0, 0, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Tetragonal, 'P'
        },
        {
            12, 2,
            true, false, 0, 0, -A / 2, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Hexagonal, 'P'
        },
        {
            13, 2,
            true, false, 0, 0, F, false, false,
            {1, 1, 0, -1, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            15, 2,
            true, false, -A / 2, -A / 2, 0, false, false,
            {1, 0, 0, 0, 1, 0, 1, 1, 2},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            16, 2,
            true, false, D, D, F, true, false,
            {-1, -1, 0, 1, -1, 0, 1, 1, 2},
            gemmi::CrystalSystem::Orthorhombic, 'F'
        },
        {
            14, 2,
            true, false, D, D, F, false, false,
            {1, 1, 0, -1, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            17, 2,
            true, false, D, E, F, true, false,
            {1, -1, 0, 1, 1, 0, -1, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            18, 1,
            false, true, A / 4, A / 2, A / 2, false, false,
            {0, -1, 1, 1, -1, -1, 1, 0, 0},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            19, 1,
            false, true, D, A / 2, A / 2, false, false,
            {-1, 0, 0, 0, -1, 1, -1, 1, 1},
            gemmi::CrystalSystem::Orthorhombic, 'I'
        },
        {
            20, 1,
            false, true, D, E, E, false, false,
            {0, 1, 1, 0, 1, -1, -1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            21, 2,
            false, true, 0, 0, 0, false, false,
            {0, 1, 0, 0, 0, 1, 1, 0, 0},
            gemmi::CrystalSystem::Tetragonal, 'P'
        },
        {
            22, 2,
            false, true, -B / 2, 0, 0, false, false,
            {0, 1, 0, 0, 0, 1, 1, 0, 0},
            gemmi::CrystalSystem::Hexagonal, 'P'
        },
        {
            23, 2,
            false, true, D, 0, 0, false, false,
            {0, 1, 1, 0, -1, 1, 1, 0, 0},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            24, 2,
            false, true, D, -A / 3, -A / 3, true, false,
            {1, 2, 1, 0, -1, 1, 1, 0, 0},
            gemmi::CrystalSystem::Trigonal, 'R'
        },
        {
            25, 2,
            false, true, D, E, E, false, false,
            {0, 1, 1, 0, -1, 1, 1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            26, 1,
            false, false, A / 4, A / 2, A / 2, false, false,
            {1, 0, 0, -1, 2, 0, -1, 0, 2},
            gemmi::CrystalSystem::Orthorhombic, 'F'
        },
        {
            27, 1,
            false, false, D, A / 2, A / 2, false, false,
            {-1, 2, 0, -1, 0, 0, 0, -1, 1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            28, 1,
            false, false, D, A / 2, 2 * D, false, false,
            {-1, 0, 0, -1, 0, 2, 0, 1, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            29, 1,
            false, false, D, 2 * D, A / 2, false, false,
            {1, 0, 0, 1, -2, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            30, 1,
            false, false, B / 2, E, 2 * E, false, false,
            {0, 1, 0, 0, 1, -2, -1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },

        {
            31, 1,
            false, false, D, E, F, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Triclinic, 'P'
        },

        {
            32, 2,
            false, false, 0, 0, 0, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Orthorhombic, 'P'
        },
        {
            40, 2,
            false, false, -B / 2, 0, 0, false, false,
            {0, -1, 0, 0, 1, 2, -1, 0, 0},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            35, 2,
            false, false, D, 0, 0, false, false,
            {0, -1, 0, -1, 0, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'P'
        },
        {
            36, 2,
            false, false, 0, -A / 2, 0, false, false,
            {1, 0, 0, -1, 0, -2, 0, 1, 0},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            33, 2,
            false, false, 0, E, 0, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Monoclinic, 'P'
        },
        {
            38, 2,
            false, false, 0, 0, -A / 2, false, false,
            {-1, 0, 0, 1, 2, 0, 0, 0, -1},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            34, 2,
            false, false, 0, 0, F, false, false,
            {-1, 0, 0, 0, 0, -1, 0, -1, 0},
            gemmi::CrystalSystem::Monoclinic, 'P'
        },
        {
            42, 2,
            false, false, -B / 2, -A / 2, 0, false, false,
            {-1, 0, 0, 0, -1, 0, 1, 1, 2},
            gemmi::CrystalSystem::Orthorhombic, 'I'
        },
        {
            41, 2,
           false, false, -B / 2, E, 0, false, false,
            {0, -1, -2, 0, -1, 0, -1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            37, 2,
            false, false, D, -A / 2, 0, false, false,
            {1, 0, 2, 1, 0, 0, 0, 1, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            39, 2,
            false, false, D, 0, -A / 2, false, false,
            {-1, -2, 0, -1, 0, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            // ITA character 43: the type-II reduced form of a CENTRED MONOCLINIC lattice with no
            // length equality. Both of its conditions are equalities on scalar products, so the three
            // angle tests are vacuous for it and cond_2DF is what selects it. mC, mI, mA and mF are
            // one Bravais lattice in four settings; this row names the reduced form whose
            // conventional cell comes out I-centred. Without it such a cell falls through to
            // character 44 and loses its centring outright.
            43, 2,
            false, false, D, E, F, true, true,
            {-1, 0, 0, -1, -1, -2, 0, -1, 0},
            gemmi::CrystalSystem::Monoclinic, 'I'
        },
        {
            44, 2,
            false, false, D, E, F, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Triclinic, 'P'
        }
      };
    };

    auto match = [&](const CrystalLattice &latt, double D, double E, double F)
            -> std::optional<LatticeSearchResult> {
        const auto uc_reduced = latt.GetUnitCell();
        for (const auto &c: make_classes(D, E, F)) {
            if (only_class && (c.system != only_class->first || c.centering != only_class->second))
                continue;
            if (c.type == 1 && uc_reduced.beta >= 90 - angle_tolerance )
                continue;

            bool ok = true;

            if (c.cond_AB && fabs((uc_reduced.a - uc_reduced.b) / (0.5 * (uc_reduced.a + uc_reduced.b))) > dist_tolerance)
                ok = false;
            if (c.cond_BC && fabs((uc_reduced.b - uc_reduced.c) / (0.5 * (uc_reduced.b + uc_reduced.c))) > dist_tolerance)
                ok = false;

            // A character states its scalar products as fractions of this cell's own A, B and C, so
            // the cosine it implies can come out beyond +/-1 - the character is then geometrically
            // impossible for this metric. acos gives NaN there, and every comparison with a NaN is
            // false, so an impossible condition used to read as a satisfied one.
            const double cos_alpha = c.cond_D / sqrt(B*C);
            const double cos_beta = c.cond_E / sqrt(A*C);
            const double cos_gamma = c.cond_F / sqrt(A*B);
            if (fabs(cos_alpha) > 1 || fabs(cos_beta) > 1 || fabs(cos_gamma) > 1)
                ok = false;

            double expected_alpha = acos(cos_alpha) * 180 / PI;
            double expected_beta = acos(cos_beta) * 180 / PI;
            double expected_gamma = acos(cos_gamma) * 180 / PI;

            if (fabs(expected_alpha - uc_reduced.alpha) > angle_tolerance)
                ok = false;
            if (fabs(expected_beta - uc_reduced.beta) > angle_tolerance)
                ok = false;
            if (fabs(expected_gamma - uc_reduced.gamma) > angle_tolerance)
                ok = false;

            double tmp1 = 2.0 * fabs(D + E + F);
            double tmp2 = A + B;

            if (c.cond_DEF && fabs((tmp1 - tmp2) / (0.5 * (tmp1 + tmp2))) > dist_tolerance)
                ok = false;

            // The second equality character 43 is made of: |2D + F| = B. Until that row existed no
            // character used cond_2DF and the test was never written.
            const double tmp3 = fabs(2.0 * D + F);
            if (c.cond_2DF && fabs((tmp3 - B) / (0.5 * (tmp3 + B))) > dist_tolerance)
                ok = false;

            if (ok) {
                return LatticeSearchResult{
                    .niggli_class = c.number,
                    .primitive_reduced = latt,
                    .conventional = latt.Multiply(c.reindex),
                    .system = c.system,
                    .centering = c.centering,
                        .reindex = c.reindex,
                };
            }
        }
        return std::nullopt;
    };

    // Character 44 fits any cell, so a match is always found - "nothing fits" is reported as triclinic.
    auto found = match(L_niggli, D, E, F);
    if (found && found->system != gemmi::CrystalSystem::Triclinic)
        return *found;

    // A reduced cell with an angle 90 to within NIGGLI_TYPE_BOUNDARY_DEG sits ON the boundary between the
    // two Niggli types: the same lattice reduces to an all-acute cell or an all-obtuse one according to
    // the last digits of whatever refinement produced it. The type-1 characters are skipped for such a
    // cell (just above) and the type-2 ones are stated for the obtuse setting, so an acute cell can match
    // none of them and comes back triclinic. Present it in the obtuse setting and try once more -
    // negating two of the three basis vectors keeps the lattice and the angle between those two and turns
    // the other two angles into their supplements. Each of alpha, beta and gamma therefore has its own
    // flip, and a cell sitting on the boundary in alpha or in gamma is not reached by the beta one.
    // Measured on a C-centred monoclinic crystal whose reduced beta sits 0.07 deg from 90 (its centring
    // was read or missed according to the sign of that 0.07 deg, and with it the space group of the whole
    // run) and on an I-centred orthorhombic one whose reduced gamma sits 0.1 deg from 90. Beta is tried
    // first, so a cell the earlier beta-only retry already rescued is answered exactly as before.
    const UnitCell uc_niggli = L_niggli.GetUnitCell();
    if (D > 0 && E > 0 && F > 0) {
        struct Flip { double kept_angle; gemmi::Mat33 basis; double d, e, f; };
        const Flip flips[3] = {
            {uc_niggli.beta,  gemmi::Mat33(-1, 0, 0, 0, 1, 0, 0, 0, -1), -D,  E, -F},
            {uc_niggli.alpha, gemmi::Mat33(1, 0, 0, 0, -1, 0, 0, 0, -1),  D, -E, -F},
            {uc_niggli.gamma, gemmi::Mat33(-1, 0, 0, 0, -1, 0, 0, 0, 1), -D, -E,  F},
        };
        for (const auto &flip : flips) {
            if (flip.kept_angle < 90 - NIGGLI_TYPE_BOUNDARY_DEG)
                continue;
            const auto flipped = match(L_niggli.Multiply(flip.basis), flip.d, flip.e, flip.f);
            if (flipped && flipped->system != gemmi::CrystalSystem::Triclinic)
                return *flipped;
        }
    }

    if (found)
        return *found;

    if (only_class)
        return std::nullopt;      // no character of the requested class fits this metric

    return LatticeSearchResult{
        .niggli_class = 44,
        .primitive_reduced = L_niggli,
        .conventional = L_niggli,
        .system = gemmi::CrystalSystem::Triclinic,
        .centering = 'P',
        .reindex = gemmi::Mat33(1, 0, 0, 0, 1, 0, 0, 0, 1),
    };
}
}   // namespace

LatticeSearchResult LatticeSearch(const CrystalLattice &L, double dist_tolerance, double angle_tolerance) {
    return *SearchCharacters(L, dist_tolerance, angle_tolerance, nullptr);
}

std::optional<LatticeSearchResult> LatticeSearchForClass(const CrystalLattice &L,
                                                         gemmi::CrystalSystem system, char centering,
                                                         double dist_tolerance, double angle_tolerance) {
    const std::pair<gemmi::CrystalSystem, char> only{system, centering};
    return SearchCharacters(L, dist_tolerance, angle_tolerance, &only);
}

double MetricDeviation(const LatticeSearchResult &sr) {
    const UnitCell uc = sr.conventional.GetUnitCell();
    double angle = 0.0, length = 0.0;
    auto a_dev = [&](double x, double ideal) { angle = std::max(angle, std::fabs(x - ideal)); };
    auto l_dev = [&](double x, double y) { length = std::max(length, std::fabs(x - y) / (0.5 * (x + y))); };
    switch (sr.system) {
        case gemmi::CrystalSystem::Monoclinic:   a_dev(uc.alpha, 90); a_dev(uc.gamma, 90); break;
        case gemmi::CrystalSystem::Orthorhombic: a_dev(uc.alpha, 90); a_dev(uc.beta, 90); a_dev(uc.gamma, 90); break;
        case gemmi::CrystalSystem::Tetragonal:   a_dev(uc.alpha, 90); a_dev(uc.beta, 90); a_dev(uc.gamma, 90);
                                                 l_dev(uc.a, uc.b); break;
        case gemmi::CrystalSystem::Trigonal:
        case gemmi::CrystalSystem::Hexagonal:    a_dev(uc.alpha, 90); a_dev(uc.beta, 90); a_dev(uc.gamma, 120);
                                                 l_dev(uc.a, uc.b); break;
        case gemmi::CrystalSystem::Cubic:        a_dev(uc.alpha, 90); a_dev(uc.beta, 90); a_dev(uc.gamma, 90);
                                                 l_dev(uc.a, uc.b); l_dev(uc.b, uc.c); break;
        default: break;                          // triclinic: nothing is asserted, nothing is violated
    }
    return std::max(angle / LATTICE_SEARCH_ANGLE_TOLERANCE_DEG, length / LATTICE_SEARCH_DIST_TOLERANCE);
}

CrystalLattice SymmetrizeMetric(const CrystalLattice &L, const gemmi::SpaceGroup &sg) {
    const Coord v[3] = {L.Vec0(), L.Vec1(), L.Vec2()};
    Eigen::Matrix3d A;
    for (int i = 0; i < 3; i++)
        A.row(i) << v[i].x, v[i].y, v[i].z;

    const Eigen::Matrix3d G = A * A.transpose();
    // The rotations alone: a centring translation acts on the lattice points, not on the metric.
    const auto &ops = sg.operations().sym_ops;
    Eigen::Matrix3d G_sym = Eigen::Matrix3d::Zero();
    for (const gemmi::Op &op : ops) {
        Eigen::Matrix3d R;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                R(i, j) = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
        G_sym += R.transpose() * G * R;
    }
    G_sym /= static_cast<double>(ops.size());

    // Symmetric square root of a positive-definite metric, or its inverse.
    const auto root = [](const Eigen::Matrix3d &M, bool inverse) {
        const Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> es(M);
        Eigen::Vector3d d = es.eigenvalues().cwiseSqrt();
        if (inverse)
            d = d.cwiseInverse();
        return Eigen::Matrix3d(es.eigenvectors() * d.asDiagonal() * es.eigenvectors().transpose());
    };
    // As a change of basis rather than three rebuilt vectors: the three-Coord constructor enforces a
    // right-handed basis, and a lattice that arrives left-handed would come back with one axis flipped
    // and its free angle replaced by the supplement (beta 132.23 -> 47.77, measured). A projection of
    // the metric has no business changing the hand of the cell its reflections are indexed on.
    const Eigen::Matrix3d M = root(G_sym, false) * root(G, true);
    return L.Multiply(gemmi::Mat33(M(0, 0), M(0, 1), M(0, 2),
                                   M(1, 0), M(1, 1), M(1, 2),
                                   M(2, 0), M(2, 1), M(2, 2)));
}