use alloc::format;
use alloc::vec;
use alloc::vec::Vec;
use core::cell::{Cell, RefCell};
use deep_causality_num::FromPrimitive;
use deep_causality_tensor::CausalTensor;
use deep_causality_topology::{
ChainComplex, DecStencilTables, HodgeDecomposeOptions, LatticeComplex, LerayProjection,
Manifold,
};
use crate::solvers::dec::DecNsScalar;
use crate::solvers::dec::spectral_diffusion::SpectralDiffusion;
use deep_causality_physics::BodyForceOneForm;
use deep_causality_physics::PhysicsError;
use deep_causality_physics::VelocityOneForm;
#[derive(Debug)]
pub struct DecNsRate<'m, const D: usize, R: DecNsScalar> {
manifold: &'m Manifold<LatticeComplex<D, R>, R>,
nu: Cell<R>,
body_force: Option<CausalTensor<R>>,
n1: usize,
engine: Option<StencilEngine<R>>,
spectral: Option<SpectralDiffusion<R>>,
no_slip: super::dec_ns_solver::no_slip::NoSlipConstraint<R>,
inflow_edges: alloc::vec::Vec<usize>,
reference_vertices: alloc::vec::Vec<usize>,
zone_constrained: alloc::vec::Vec<usize>,
rate_constrained: alloc::vec::Vec<usize>,
warm_start: bool,
proj_warm: RefCell<Option<alloc::vec::Vec<R>>>,
proj_warm_lambda: RefCell<Option<alloc::vec::Vec<R>>>,
}
#[derive(Debug)]
struct StencilEngine<R> {
tables: DecStencilTables<R>,
ws: RefCell<RateWorkspace<R>>,
}
#[derive(Debug)]
struct RateWorkspace<R> {
omega: Vec<R>,
pre: Vec<R>,
wedge: Vec<R>,
conv: Vec<R>,
visc_a: Vec<R>,
visc_b: Vec<R>,
s0: Vec<R>,
adj_corr: Vec<R>,
adj_s1: Vec<R>,
adj_sw: Vec<R>,
}
impl<'m, const D: usize, R: DecNsScalar> DecNsRate<'m, D, R> {
pub fn new(
manifold: &'m Manifold<LatticeComplex<D, R>, R>,
nu: R,
body_force: Option<&BodyForceOneForm<R>>,
) -> Result<Self, PhysicsError> {
if D < 2 {
return Err(PhysicsError::DimensionMismatch(format!(
"DecNsRate requires a lattice of dimension >= 2 (the convective \
term contracts a grade-2 vorticity), got D = {D}"
)));
}
if manifold.metric().is_none() {
return Err(PhysicsError::TopologyError(
"DecNsRate requires a metric-bearing manifold (Hodge star); \
construct it with CubicalReggeGeometry"
.into(),
));
}
if !nu.is_finite() {
return Err(PhysicsError::NumericalInstability(
"DecNsRate: viscosity must be finite".into(),
));
}
if nu < R::zero() {
return Err(PhysicsError::PhysicalInvariantBroken(
"DecNsRate: viscosity cannot be negative".into(),
));
}
let complex = manifold.complex();
let n1 = complex.num_cells(1);
let cut_registry = manifold.metric().and_then(|m| m.cut_registry());
let no_slip =
super::dec_ns_solver::no_slip::NoSlipConstraint::new(complex, cut_registry, true);
let any_wall = complex.periodic().iter().any(|&p| !p);
if any_wall {
for (axis, (&periodic, &extent)) in complex
.periodic()
.iter()
.zip(complex.shape().iter())
.enumerate()
{
if !periodic && extent < 2 {
return Err(PhysicsError::DimensionMismatch(format!(
"DecNsRate: wall axis {axis} has extent {extent}; wall-bounded \
lattices need at least 2 vertex layers per wall axis"
)));
}
}
}
if any_wall || cut_registry.is_some() {
use deep_causality_topology::HasHodgeStar;
let immersed: alloc::vec::Vec<usize> = cut_registry
.map(|r| r.solid_incident_edges(complex))
.unwrap_or_default();
let metric = manifold
.metric()
.expect("metric presence checked above");
let star = metric
.hodge_star_matrix(complex, 1)
.map_err(|e| PhysicsError::TopologyError(format!("hodge star (grade 1): {e}")))?;
for i in 0..n1 {
if immersed.binary_search(&i).is_ok() {
continue;
}
let mut diag = R::zero();
for e in star.row_indices()[i]..star.row_indices()[i + 1] {
if star.col_indices()[e] == i {
diag = star.values()[e];
}
}
if !diag.is_finite() || diag <= R::zero() {
return Err(PhysicsError::TopologyError(format!(
"DecNsRate: free edges require the boundary-corrected Hodge star with \
strictly positive masses; edge {i} has mass {diag}"
)));
}
}
}
let body_force = match body_force {
Some(g) => {
if g.len() != n1 {
return Err(PhysicsError::DimensionMismatch(format!(
"DecNsRate: body force carries {} edge coefficients, \
the lattice has {n1}",
g.len()
)));
}
Some(g.as_tensor().clone())
}
None => None,
};
let tables = DecStencilTables::compile(manifold)
.map_err(|e| PhysicsError::TopologyError(format!("stencil compilation failed: {e}")))?;
let n0 = complex.num_cells(0);
let n2 = complex.num_cells(2);
let (pre_len, wedge_len) = tables.convective_scratch_lens();
let (adj_s1_len, adj_sw_len) = tables.convective_vector_adjoint_scratch_lens();
let ws = RateWorkspace {
omega: vec![R::zero(); n2],
pre: vec![R::zero(); pre_len],
wedge: vec![R::zero(); wedge_len],
conv: vec![R::zero(); n1],
visc_a: vec![R::zero(); n1],
visc_b: vec![R::zero(); n1],
s0: vec![R::zero(); n0],
adj_corr: vec![R::zero(); n1],
adj_s1: vec![R::zero(); adj_s1_len],
adj_sw: vec![R::zero(); adj_sw_len],
};
let engine = Some(StencilEngine {
tables,
ws: RefCell::new(ws),
});
let rate_constrained = no_slip.edges().to_vec();
Ok(Self {
manifold,
nu: Cell::new(nu),
body_force,
n1,
engine,
spectral: None,
no_slip,
inflow_edges: alloc::vec::Vec::new(),
reference_vertices: alloc::vec::Vec::new(),
zone_constrained: alloc::vec::Vec::new(),
rate_constrained,
warm_start: false,
proj_warm: RefCell::new(None),
proj_warm_lambda: RefCell::new(None),
})
}
pub(in crate::solvers::dec) fn set_warm_start(&mut self, on: bool) {
self.warm_start = on;
if !on {
*self.proj_warm.borrow_mut() = None;
*self.proj_warm_lambda.borrow_mut() = None;
}
}
pub(in crate::solvers::dec) fn set_open_boundary(
&mut self,
inflow: alloc::vec::Vec<usize>,
reference: alloc::vec::Vec<usize>,
) {
let mut inflow = inflow;
inflow.sort_unstable();
inflow.dedup();
let mut reference = reference;
reference.sort_unstable();
reference.dedup();
self.inflow_edges = inflow;
self.reference_vertices = reference;
self.recompute_rate_constrained();
}
pub(in crate::solvers::dec) fn apply_slip(&mut self, slip: &[usize]) {
if slip.is_empty() {
return;
}
self.no_slip.remove_edges(slip);
self.recompute_rate_constrained();
}
pub(in crate::solvers::dec) fn set_staircase_noslip(&mut self) {
let complex = self.manifold.complex();
let cut_registry = self.manifold.metric().and_then(|m| m.cut_registry());
self.no_slip =
super::dec_ns_solver::no_slip::NoSlipConstraint::new(complex, cut_registry, false);
self.recompute_rate_constrained();
}
pub(in crate::solvers::dec) fn set_zone_constrained(&mut self, edges: alloc::vec::Vec<usize>) {
let mut edges = edges;
edges.sort_unstable();
edges.dedup();
self.zone_constrained = edges;
self.recompute_rate_constrained();
}
fn recompute_rate_constrained(&mut self) {
let mut rate_constrained = self.no_slip.edges().to_vec();
rate_constrained.extend_from_slice(&self.inflow_edges);
rate_constrained.extend_from_slice(&self.zone_constrained);
rate_constrained.sort_unstable();
rate_constrained.dedup();
self.rate_constrained = rate_constrained;
}
pub(in crate::solvers::dec) fn inflow_edges(&self) -> &[usize] {
&self.inflow_edges
}
pub(in crate::solvers::dec) fn reference_vertices(&self) -> &[usize] {
&self.reference_vertices
}
pub(in crate::solvers::dec) fn no_slip_edges(&self) -> &[usize] {
self.no_slip.edges()
}
pub(in crate::solvers::dec) fn no_slip_rows(
&self,
) -> &[deep_causality_topology::CutFaceConstraint<R>] {
self.no_slip.rows()
}
pub fn with_spectral_diffusion(mut self) -> Result<Self, PhysicsError> {
self.spectral = Some(SpectralDiffusion::new(self.manifold)?);
Ok(self)
}
pub fn with_generic_assembly(mut self) -> Self {
self.engine = None;
self
}
pub fn nu(&self) -> R {
self.nu.get()
}
pub fn set_nu(&self, nu: R) {
self.nu.set(nu);
}
pub fn eval_projected(
&self,
u: &VelocityOneForm<R>,
opts: &HodgeDecomposeOptions<R>,
) -> Result<VelocityOneForm<R>, PhysicsError> {
let raw = self.eval_unprojected(u);
let projection = self.project_raw(&raw, opts)?;
let (projected, _potential) = projection.into_parts();
Ok(VelocityOneForm::from_raw(projected))
}
pub(crate) fn eval_projected_with_potential(
&self,
u: &VelocityOneForm<R>,
opts: &HodgeDecomposeOptions<R>,
) -> Result<(VelocityOneForm<R>, CausalTensor<R>), PhysicsError> {
let raw = self.eval_unprojected(u);
let projection = self.project_raw(&raw, opts)?;
let (projected, potential) = projection.into_parts();
Ok((VelocityOneForm::from_raw(projected), potential))
}
pub fn energy_budget(
&self,
u: &VelocityOneForm<R>,
opts: &HodgeDecomposeOptions<R>,
) -> Result<super::energy_budget::EnergyBudget<R>, PhysicsError> {
let u_slice = u.as_tensor().as_slice();
let (conv, lap): (Vec<R>, Vec<R>) = if let Some(engine) = &self.engine {
let t = &engine.tables;
let mut ws = engine.ws.borrow_mut();
let ws = &mut *ws;
t.apply_d1(u_slice, &mut ws.omega)
.expect("workspace lengths fixed at construction");
Self::fill_convective_skew_fused(t, ws, u_slice);
if let Some(spectral) = &self.spectral {
spectral.apply_laplacian_1(u_slice, &mut ws.visc_a);
for v in ws.visc_b.iter_mut() {
*v = R::zero();
}
} else {
t.apply_delta2(&ws.omega, &mut ws.visc_a)
.expect("workspace lengths fixed at construction");
t.apply_delta1(u_slice, &mut ws.s0)
.expect("workspace lengths fixed at construction");
t.apply_d0(&ws.s0, &mut ws.visc_b)
.expect("workspace lengths fixed at construction");
}
let lap = ws
.visc_a
.iter()
.zip(ws.visc_b.iter())
.map(|(a, b)| *a + *b)
.collect();
(ws.conv.clone(), lap)
} else {
let conv = self.convective_skew_generic(u);
let mut lap = self.manifold.laplacian_of(u_slice, 1).into_vec();
lap.resize(self.n1, R::zero());
(conv, lap)
};
let m_inner = |v: &[R]| -> R {
let star_v = self.manifold.hodge_star_of(v, 1);
u_slice
.iter()
.zip(star_v.as_slice().iter())
.fold(R::zero(), |acc, (a, b)| acc + *a * *b)
};
let convective = R::zero() - m_inner(&conv);
let viscous = R::zero() - self.nu.get() * m_inner(&lap);
let body_force = match &self.body_force {
Some(g) => m_inner(g.as_slice()),
None => R::zero(),
};
let projected_rate = self.eval_projected(u, opts)?;
let projected = m_inner(projected_rate.as_tensor().as_slice());
Ok(super::energy_budget::EnergyBudget {
convective,
viscous,
body_force,
projected,
})
}
fn project_raw(
&self,
raw: &VelocityOneForm<R>,
opts: &HodgeDecomposeOptions<R>,
) -> Result<LerayProjection<R>, PhysicsError> {
let rows = self.no_slip.rows();
if !self.warm_start {
return self
.manifold
.leray_project_constrained_weighted_opts(
raw.as_tensor(),
&self.rate_constrained,
rows,
opts,
None,
)
.map_err(|e| PhysicsError::TopologyError(format!("Leray projection failed: {e}")));
}
let (projection, lambda) = {
let phi_guess = self.proj_warm.borrow();
let lambda_guess = self.proj_warm_lambda.borrow();
self.manifold
.leray_project_constrained_weighted_warm(
raw.as_tensor(),
&self.rate_constrained,
rows,
opts,
phi_guess.as_deref(),
lambda_guess.as_deref(),
)
.map_err(|e| PhysicsError::TopologyError(format!("Leray projection failed: {e}")))?
};
*self.proj_warm.borrow_mut() = Some(projection.potential().as_slice().to_vec());
*self.proj_warm_lambda.borrow_mut() = Some(lambda);
Ok(projection)
}
fn fill_convective_skew_fused(
tables: &DecStencilTables<R>,
ws: &mut RateWorkspace<R>,
u_slice: &[R],
) {
tables
.apply_convective(&ws.omega, u_slice, &mut ws.pre, &mut ws.wedge, &mut ws.conv)
.expect("workspace lengths fixed at construction");
tables
.apply_convective_vector_adjoint(
&ws.pre,
u_slice,
&mut ws.adj_s1,
&mut ws.adj_sw,
&mut ws.adj_corr,
)
.expect("workspace lengths fixed at construction");
let half = R::from_f64(0.5)
.expect("0.5 lifts into R");
for (c, k) in ws.conv.iter_mut().zip(ws.adj_corr.iter()) {
*c = half * (*c - *k);
}
}
fn convective_skew_generic(&self, u: &VelocityOneForm<R>) -> Vec<R> {
let u_slice = u.as_tensor().as_slice();
let du = self.manifold.exterior_derivative_of(u_slice, 1);
let conv_raw = self
.manifold
.interior_product(u.as_tensor(), &du, 2)
.expect("interior_product preconditions validated at construction")
.into_vec();
let m1 = self.manifold.hodge_star_of(&vec![R::one(); self.n1], 1);
let m1 = m1.as_slice();
let zero_tol = <R as FromPrimitive>::from_f64(1e-12)
.expect("1e-12 is representable in every RealField");
let w: Vec<R> = u_slice
.iter()
.zip(m1.iter())
.map(|(a, b)| *a * *b)
.collect();
let mut adj = vec![R::zero(); self.n1];
for (j, slot) in adj.iter_mut().enumerate() {
let mut e = vec![R::zero(); self.n1];
e[j] = R::one();
let e_t = CausalTensor::new(e, vec![self.n1])
.expect("1-D tensor allocation cannot fail");
let col = self
.manifold
.interior_product(&e_t, &du, 2)
.expect("interior_product preconditions validated at construction");
let dot = col
.as_slice()
.iter()
.zip(w.iter())
.fold(R::zero(), |acc, (a, b)| acc + *a * *b);
*slot = if m1[j].abs() <= zero_tol {
R::zero()
} else {
dot / m1[j]
};
}
let half = R::from_f64(0.5)
.expect("0.5 lifts into R");
conv_raw
.iter()
.zip(adj.iter())
.map(|(c, k)| half * (*c - *k))
.collect()
}
pub fn eval_unprojected(&self, u: &VelocityOneForm<R>) -> VelocityOneForm<R> {
debug_assert_eq!(
u.len(),
self.n1,
"marching state length is invariant under Add/Mul and validated at seeding"
);
if let Some(engine) = &self.engine {
return self.eval_unprojected_fused(engine, u);
}
let u_slice = u.as_tensor().as_slice();
let conv = self.convective_skew_generic(u);
let lap = self.manifold.laplacian_of(u_slice, 1);
let conv_s = conv.as_slice();
let lap_s = lap.as_slice();
let rhs: Vec<R> = match &self.body_force {
Some(g) => {
let g_s = g.as_slice();
let nu = self.nu.get();
(0..self.n1)
.map(|i| R::zero() - conv_s[i] - nu * lap_s[i] + g_s[i])
.collect()
}
None => {
let nu = self.nu.get();
(0..self.n1)
.map(|i| R::zero() - conv_s[i] - nu * lap_s[i])
.collect()
}
};
let tensor = CausalTensor::new(rhs, vec![self.n1])
.expect("1-D tensor allocation cannot fail");
VelocityOneForm::from_raw(tensor)
}
fn eval_unprojected_fused(
&self,
engine: &StencilEngine<R>,
u: &VelocityOneForm<R>,
) -> VelocityOneForm<R> {
let u_slice = u.as_tensor().as_slice();
let t = &engine.tables;
let mut ws = engine.ws.borrow_mut();
let ws = &mut *ws;
t.apply_d1(u_slice, &mut ws.omega)
.expect("workspace lengths fixed at construction");
Self::fill_convective_skew_fused(t, ws, u_slice);
if let Some(spectral) = &self.spectral {
spectral.apply_laplacian_1(u_slice, &mut ws.visc_a);
for v in ws.visc_b.iter_mut() {
*v = R::zero();
}
} else {
t.apply_delta2(&ws.omega, &mut ws.visc_a)
.expect("workspace lengths fixed at construction");
t.apply_delta1(u_slice, &mut ws.s0)
.expect("workspace lengths fixed at construction");
t.apply_d0(&ws.s0, &mut ws.visc_b)
.expect("workspace lengths fixed at construction");
}
let rhs: Vec<R> = match &self.body_force {
Some(g) => {
let g_s = g.as_slice();
let nu = self.nu.get();
(0..self.n1)
.map(|i| R::zero() - ws.conv[i] - nu * (ws.visc_a[i] + ws.visc_b[i]) + g_s[i])
.collect()
}
None => {
let nu = self.nu.get();
(0..self.n1)
.map(|i| R::zero() - ws.conv[i] - nu * (ws.visc_a[i] + ws.visc_b[i]))
.collect()
}
};
let tensor = CausalTensor::new(rhs, vec![self.n1])
.expect("1-D tensor allocation cannot fail");
VelocityOneForm::from_raw(tensor)
}
}