From 0ceff6e91bdd5c042f76bd13de0f4d8c4b13968b Mon Sep 17 00:00:00 2001
From: leonarski_f <filip.leonarski@psi.ch>
Date: Sun, 17 May 2026 12:33:10 +0200
Subject: [PATCH] HDF5NXmx: Save unit cell after rotation indexing as main one
 + lattice (in UB-matrix format) + input unit cell as extra information

---
 common/CrystalLattice.cpp | 12 ++++++++++++
 common/CrystalLattice.h   |  1 +
 writer/HDF5NXmx.cpp       | 30 ++++++++++++++++++++++++++----
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/common/CrystalLattice.cpp b/common/CrystalLattice.cpp
index 32ce9350..968d28f4 100644
--- a/common/CrystalLattice.cpp
+++ b/common/CrystalLattice.cpp
@@ -199,3 +199,15 @@ void CrystalLattice::Regularize(const gemmi::CrystalSystem &input) {
             break;
     }
 }
+
+std::vector<float> CrystalLattice::GetUBMatrix() const {
+    const Coord astar = Astar();
+    const Coord bstar = Bstar();
+    const Coord cstar = Cstar();
+
+    return {
+        astar.x, bstar.x, cstar.x,
+        astar.y, bstar.y, cstar.y,
+        astar.z, bstar.z, cstar.z
+    };
+}
diff --git a/common/CrystalLattice.h b/common/CrystalLattice.h
index c7b8d358..7ec162d7 100644
--- a/common/CrystalLattice.h
+++ b/common/CrystalLattice.h
@@ -38,6 +38,7 @@ public:
     CrystalLattice Multiply(const gemmi::Mat33 &input) const;
     void Sort();
     void Regularize(const gemmi::CrystalSystem &input);
+    [[nodiscard]] std::vector<float> GetUBMatrix() const;
 };
 
 
diff --git a/writer/HDF5NXmx.cpp b/writer/HDF5NXmx.cpp
index f441b27c..a570f1db 100644
--- a/writer/HDF5NXmx.cpp
+++ b/writer/HDF5NXmx.cpp
@@ -542,6 +542,12 @@ void NXmx::Metrology(const StartMessage &start) {
                    start.pixel_size_x);
 }
 
+void SaveUnitCell( HDF5Group& group, const std::string& name, const UnitCell& unit_cell,) {
+    std::vector<float> v = {unit_cell.a, unit_cell.b, unit_cell.c,
+                            unit_cell.alpha, unit_cell.beta, unit_cell.gamma};
+    group.SaveVector(name, v);
+}
+
 void NXmx::Sample(const StartMessage &start, const EndMessage &end) {
     HDF5Group group(*hdf5_file, "/entry/sample");
     group.NXClass("NXsample");
@@ -555,10 +561,26 @@ void NXmx::Sample(const StartMessage &start, const EndMessage &end) {
             group.SaveScalar("space_group", sg->short_name());
     }
 
-    if (start.unit_cell) {
-        std::vector<float> v = {start.unit_cell->a, start.unit_cell->b, start.unit_cell->c,
-                                start.unit_cell->alpha, start.unit_cell->beta, start.unit_cell->gamma};
-        group.SaveVector("unit_cell", v);
+    std::optional<UnitCell> unit_cell;
+    std::optional<UnitCell> set_unit_cell;
+    if (end.rotation_lattice)
+        unit_cell = end.rotation_lattice->GetUnitCell();
+    else if (start.unit_cell) {
+        unit_cell = start.unit_cell;
+        set_unit_cell = start.unit_cell;
+    }
+    
+    if (unit_cell)
+        SaveUnitCell(group, "unit_cell", unit_cell.value());
+
+    if (set_unit_cell)
+        SaveUnitCell(group, "input_unit_cell", unit_cell.value());
+
+    if (end.rotation_lattice) {
+        group.SaveVector("ub_matrix",
+                         end.rotation_lattice->GetUBMatrix(),
+                         {1, 3, 3})
+                ->Units("Angstrom^-1");
     }
 
     if (start.sample_temperature_K)