gistfile1.txt

/*
I have a 1200x600 matrix tensor in open3d c++ for reflectivity and one for azimuth.

In Open3d using the GPU, find all points where azimut crosses from negative to positive. There will be one for each row in the matrix. then select 590 reflectivity values left and right of this point but in the reflectivity matrix. These slices of reflectivity should be copied into a new matrix reflectivity_adjusted that will be 1180x600 size. 


*/


#include <open3d/core/Tensor.h>
#include <open3d/core/TensorKey.h>

namespace o3c = open3d::core;

// reflectivity: float tensor [600, 1200] on CUDA
// azimuth:      float tensor [600, 1200] on CUDA
o3c::Tensor AdjustReflectivityAroundAzimuthCrossing(
        const o3c::Tensor& reflectivity,
        const o3c::Tensor& azimuth) {
    const auto device = azimuth.GetDevice();  // should be CUDA
    const int64_t H = 600;
    const int64_t W = 1200;
    const int64_t HALF = 590;
    const int64_t OUTW = 2 * HALF;  // 1180

    // --- 1) Find crossing per row: azimuth crosses from negative to positive ---
    // mask[i,j] true when az[i,j] < 0 and az[i,j+1] >= 0
    o3c::Tensor az_left  = azimuth.Slice(1, 0, W - 1);  // [H, 1199]
    o3c::Tensor az_right = azimuth.Slice(1, 1, W);      // [H, 1199]

    o3c::Tensor mask = az_left.LT(0).LogicalAnd(az_right.GE(0));  // bool [H,1199]

    // ArgMax along columns gives index of the "true" position (tie -> first).
    // Crossing column is (idx + 1) because mask is between j and j+1.
    o3c::Tensor cross_col = mask.To(o3c::Int64).ArgMax(/*axis=*/1) + 1; // [H]

    // Handle wrap-around crossing between last and first column:
    // if az[-1] < 0 and az[0] >= 0 => crossing at column 0
    o3c::Tensor last_col = azimuth.Slice(1, W - 1, W).Reshape({H}); // [H]
    o3c::Tensor first_col = azimuth.Slice(1, 0, 1).Reshape({H});    // [H]
    o3c::Tensor wrap_mask = last_col.LT(0).LogicalAnd(first_col.GE(0)); // bool [H]

    cross_col = o3c::Where(
        wrap_mask,
        o3c::Tensor::Zeros({H}, o3c::Int64, device),
        cross_col
    ); // [H] final crossing column per row

    // --- 2) Build per-row column indices for the 1180 window, with modulo wrap ---
    o3c::Tensor start = (cross_col - HALF).Reshape({H, 1}); // [H,1]

    o3c::Tensor offs = o3c::Tensor::Arange(0, OUTW, 1, o3c::Int64, device).Reshape({1, OUTW}); // [1,1180]

    // cols = (start + offs) mod 1200
    o3c::Tensor cols = (start + offs).Remainder(W); // [H,1180]

    // rows grid [H,1180]
    o3c::Tensor rows = o3c::Tensor::Arange(0, H, 1, o3c::Int64, device).Reshape({H, 1}).Broadcast({H, OUTW});

    // Linear indices into flattened reflectivity
    o3c::Tensor lin = rows * W + cols; // [H,1180]

    // --- 3) Gather reflectivity using linear indexing ---
    o3c::Tensor flat = reflectivity.Reshape({H * W}); // [720000]

    o3c::Tensor gathered_flat = flat.IndexGet({
        o3c::TensorKey::IndexTensor(lin.Reshape({H * OUTW}))
    });

    o3c::Tensor reflectivity_adjusted_hw = gathered_flat.Reshape({H, OUTW}); // [600,1180]

    // You asked for 1180x600
    o3c::Tensor reflectivity_adjusted = reflectivity_adjusted_hw.Transpose({1, 0}); // [1180,600]
    return reflectivity_adjusted;
}