use crate::PhysicsError;
use deep_causality_num::{Field, Float};
use deep_causality_tensor::CausalTensor;
#[inline]
pub fn pair_to_lie_index(mu: usize, nu: usize) -> Option<usize> {
if mu >= nu || mu >= 4 || nu >= 4 {
return None;
}
let idx = match (mu, nu) {
(0, 1) => 0,
(0, 2) => 1,
(0, 3) => 2,
(1, 2) => 3,
(1, 3) => 4,
(2, 3) => 5,
_ => return None,
};
Some(idx)
}
#[inline]
pub fn lie_index_to_pair(lie_idx: usize) -> Option<(usize, usize)> {
match lie_idx {
0 => Some((0, 1)),
1 => Some((0, 2)),
2 => Some((0, 3)),
3 => Some((1, 2)),
4 => Some((1, 3)),
5 => Some((2, 3)),
_ => None,
}
}
pub fn expand_lie_to_riemann<T>(lie_fs: &CausalTensor<T>) -> Result<CausalTensor<T>, PhysicsError>
where
T: Field + Float + Copy,
{
let shape = lie_fs.shape();
if shape.len() < 3 {
return Err(PhysicsError::DimensionMismatch(format!(
"Expected at least 3D tensor, got shape {:?}",
shape
)));
}
let lie_data = lie_fs.as_slice();
let dim = 4usize;
let _lie_dim = 6usize;
let (num_points, elem_per_lie_point, rho_stride, sigma_stride) = if shape.len() == 3 {
(
1,
shape[0] * shape[1] * shape[2],
shape[1] * shape[2],
shape[2],
)
} else {
let elem = shape[1] * shape[2] * shape[3];
(shape[0], elem, shape[2] * shape[3], shape[3])
};
let elem_per_riemann_point = dim * dim * dim * dim; let mut riemann = vec![T::zero(); num_points * elem_per_riemann_point];
let riemann_idx = |p: usize, rho: usize, sigma: usize, mu: usize, nu: usize| {
p * elem_per_riemann_point + ((rho * dim + sigma) * dim + mu) * dim + nu
};
for p in 0..num_points {
let point_offset = p * elem_per_lie_point;
for rho in 0..dim {
for sigma in 0..dim {
for mu in 0..dim {
for nu in 0..dim {
let value = if mu == nu {
T::zero()
} else if mu < nu {
if let Some(lie_idx) = pair_to_lie_index(mu, nu) {
let flat_idx = point_offset
+ rho * rho_stride
+ sigma * sigma_stride
+ lie_idx;
lie_data.get(flat_idx).copied().unwrap_or(T::zero())
} else {
T::zero()
}
} else {
if let Some(lie_idx) = pair_to_lie_index(nu, mu) {
let flat_idx = point_offset
+ rho * rho_stride
+ sigma * sigma_stride
+ lie_idx;
T::zero() - lie_data.get(flat_idx).copied().unwrap_or(T::zero())
} else {
T::zero()
}
};
riemann[riemann_idx(p, rho, sigma, mu, nu)] = value;
}
}
}
}
}
let output_shape = if num_points == 1 {
vec![dim, dim, dim, dim]
} else {
vec![num_points, dim, dim, dim, dim]
};
Ok(CausalTensor::from_vec(riemann, &output_shape))
}
pub fn contract_riemann_to_lie<T>(
riemann: &CausalTensor<T>,
) -> Result<CausalTensor<T>, PhysicsError>
where
T: Field + Float + Copy,
{
let shape = riemann.shape();
if shape != [4, 4, 4, 4] {
return Err(PhysicsError::DimensionMismatch(format!(
"Expected [4, 4, 4, 4] Riemann tensor, got {:?}",
shape
)));
}
let r_data = riemann.as_slice();
let dim = 4usize;
let lie_dim = 6usize;
let mut lie_fs = vec![T::zero(); dim * dim * lie_dim];
let riemann_idx = |rho: usize, sigma: usize, mu: usize, nu: usize| {
((rho * dim + sigma) * dim + mu) * dim + nu
};
for rho in 0..dim {
for sigma in 0..dim {
for lie_idx in 0..lie_dim {
if let Some((mu, nu)) = lie_index_to_pair(lie_idx) {
let r_value = r_data[riemann_idx(rho, sigma, mu, nu)];
lie_fs[rho * dim * lie_dim + sigma * lie_dim + lie_idx] = r_value;
}
}
}
}
Ok(CausalTensor::from_vec(lie_fs, &[dim, dim, lie_dim]))
}