pub(super) fn accumulate_msd(
    acc: &mut [f64],
    lag_idx: usize,
    cols: usize,
    axis: Option<[f64; 3]>,
    n_groups: usize,
    type_ids: &[usize],
    type_counts: &[usize],
    cur: &[f32],
    old: &[f32],
) {
    let n_types = type_counts.len();
    let n_groups_f = n_groups as f64;
    let base = lag_idx * cols;
    let block = n_types + 1;
    for g in 0..n_groups {
        let idx = g * 3;
        let dx = (cur[idx] - old[idx]) as f64;
        let dy = (cur[idx + 1] - old[idx + 1]) as f64;
        let dz = (cur[idx + 2] - old[idx + 2]) as f64;
        let msd_x = dx * dx;
        let msd_y = dy * dy;
        let msd_z = dz * dz;
        let msd_tot = msd_x + msd_y + msd_z;
        let msd_axis = axis.map(|a| {
            let proj = dx * a[0] + dy * a[1] + dz * a[2];
            proj * proj
        });
        let type_id = type_ids[g];
        let type_count = type_counts[type_id] as f64;
        let mut offset = base + type_id;
        acc[offset] += msd_x / type_count;
        acc[offset + block] += msd_y / type_count;
        acc[offset + 2 * block] += msd_z / type_count;
        if let Some(axis_val) = msd_axis {
            acc[offset + 3 * block] += axis_val / type_count;
            acc[offset + 4 * block] += msd_tot / type_count;
        } else {
            acc[offset + 3 * block] += msd_tot / type_count;
        }

        offset = base + n_types;
        acc[offset] += msd_x / n_groups_f;
        acc[offset + block] += msd_y / n_groups_f;
        acc[offset + 2 * block] += msd_z / n_groups_f;
        if let Some(axis_val) = msd_axis {
            acc[offset + 3 * block] += axis_val / n_groups_f;
            acc[offset + 4 * block] += msd_tot / n_groups_f;
        } else {
            acc[offset + 3 * block] += msd_tot / n_groups_f;
        }
    }
}
