use laddu_expr::{BinaryOp, ExprId, ExprNode, UnaryOp};
use nalgebra::{DMatrix, DVector};
use num::complex::Complex64;
use std::collections::HashMap;
use super::super::{
CpuBatchCache, CpuPlan, DynamicLu, RuntimeError, RuntimeResult, Value, matrix_at,
matrix_values_row_major, scalar_at, vector_at,
};
pub(in crate::cpu) struct DerivativeWorkspace<'a> {
plan: &'a CpuPlan,
primals: &'a [Value],
tangents: Vec<Option<Value>>,
factors: HashMap<usize, DynamicLu>,
cached_factors: Option<(&'a CpuBatchCache, usize)>,
}
impl<'a> DerivativeWorkspace<'a> {
pub(in crate::cpu) fn new(
plan: &'a CpuPlan,
primals: &'a [Value],
cached_factors: Option<(&'a CpuBatchCache, usize)>,
) -> Self {
Self {
plan,
primals,
tangents: vec![None; plan.graph.nodes().len()],
factors: HashMap::new(),
cached_factors,
}
}
pub(in crate::cpu) fn gradient(&mut self) -> RuntimeResult<Vec<Complex64>> {
let mut gradient = Vec::with_capacity(self.plan.autodiff.parameter_count());
for parameter in 0..self.plan.autodiff.parameter_count() {
let active = self
.plan
.autodiff
.active_nodes(parameter)
.expect("free parameter index is valid");
for id in active {
self.differentiate_node(*id)?;
}
gradient.push(self.scalar_tangent(self.plan.graph.root())?);
for id in active {
self.tangents[id.index()] = None;
}
}
Ok(gradient)
}
fn differentiate_node(&mut self, id: ExprId) -> RuntimeResult<()> {
let index = id.index();
let node = self.plan.graph.nodes()[index].clone();
let tangent = match node {
ExprNode::ScalarParam(_) => Value::Scalar(Complex64::ONE),
ExprNode::Unary { op, input } => {
let input_value = self.primal_scalar(input)?;
let output_value = self.primal_scalar(id)?;
let input_tangent = self.scalar_tangent(input)?;
let value = match op {
UnaryOp::Neg => -input_tangent,
UnaryOp::Real => Complex64::from(input_tangent.re),
UnaryOp::Imag => Complex64::from(input_tangent.im),
UnaryOp::Conj => input_tangent.conj(),
UnaryOp::NormSqr => {
Complex64::from(2.0 * (input_value.conj() * input_tangent).re)
}
UnaryOp::Sqrt => input_tangent / (2.0 * output_value),
UnaryOp::Exp => output_value * input_tangent,
UnaryOp::Sin => input_value.cos() * input_tangent,
UnaryOp::Cos => -input_value.sin() * input_tangent,
UnaryOp::Log => input_tangent / input_value,
UnaryOp::PowI(power) => {
if power == 0 {
Complex64::ZERO
} else if power == i32::MIN {
power as f64 * output_value * input_tangent / input_value
} else {
power as f64 * input_value.powi(power - 1) * input_tangent
}
}
};
Value::Scalar(value)
}
ExprNode::Binary { op, lhs, rhs } => {
let lhs_value = self.primal_scalar(lhs)?;
let rhs_value = self.primal_scalar(rhs)?;
let lhs_tangent = self.scalar_tangent(lhs)?;
let rhs_tangent = self.scalar_tangent(rhs)?;
let value = match op {
BinaryOp::Add => lhs_tangent + rhs_tangent,
BinaryOp::Sub => lhs_tangent - rhs_tangent,
BinaryOp::Mul => lhs_tangent * rhs_value + lhs_value * rhs_tangent,
BinaryOp::Div => {
(lhs_tangent * rhs_value - lhs_value * rhs_tangent) / rhs_value.powi(2)
}
BinaryOp::Atan2 => {
let denominator = lhs_value.re.powi(2) + rhs_value.re.powi(2);
Complex64::from(
(rhs_value.re * lhs_tangent.re - lhs_value.re * rhs_tangent.re)
/ denominator,
)
}
};
Value::Scalar(value)
}
ExprNode::NaryAdd { terms } => {
Value::Scalar(terms.into_iter().try_fold(Complex64::ZERO, |sum, term| {
Ok::<_, RuntimeError>(sum + self.scalar_tangent(term)?)
})?)
}
ExprNode::NaryMul { factors } => {
let mut product = Complex64::ONE;
let mut derivative = Complex64::ZERO;
for factor in factors {
let value = self.primal_scalar(factor)?;
derivative = derivative * value + product * self.scalar_tangent(factor)?;
product *= value;
}
Value::Scalar(derivative)
}
ExprNode::Complex { re, im } => Value::Scalar(Complex64::new(
self.scalar_tangent(re)?.re,
self.scalar_tangent(im)?.re,
)),
ExprNode::Vector { .. }
if self.cached_factors.is_some()
&& self.plan.cached_value_slots[index].is_none()
&& self.plan.solve_rhs_elements[index].is_some() =>
{
Value::Vector(Vec::new())
}
ExprNode::Vector { elements } => Value::Vector(
elements
.into_iter()
.map(|element| self.scalar_tangent(element))
.collect::<RuntimeResult<_>>()?,
),
ExprNode::Matrix {
rows,
cols,
elements,
} => {
if elements.len() != rows * cols {
return Err(RuntimeError::InvalidShape {
index,
message: format!(
"matrix has {} elements for shape {rows}x{cols}",
elements.len()
),
});
}
Value::Matrix {
rows,
cols,
values: elements
.into_iter()
.map(|element| self.scalar_tangent(element))
.collect::<RuntimeResult<_>>()?,
}
}
ExprNode::Component { input, index: i } => {
if let (Some(plan), Some((cache, row))) =
(self.plan.solve_components[index], self.cached_factors)
{
let inverse_row = cache.solve_row(plan.row_slot(), row)?;
if let Some(elements) = &self.plan.solve_rhs_elements[plan.rhs().index()] {
Value::Scalar(
inverse_row
.iter()
.zip(elements)
.map(|(lhs, rhs)| Ok(lhs * self.scalar_tangent(*rhs)?))
.sum::<RuntimeResult<Complex64>>()?,
)
} else {
let rhs_tangent =
self.vector_tangent_value(plan.rhs(), plan.dimension())?;
Value::Scalar(
inverse_row
.iter()
.zip(rhs_tangent)
.map(|(lhs, rhs)| lhs * rhs)
.sum(),
)
}
} else {
let vector = self.vector_tangent(input)?;
Value::Scalar(*vector.get(i).ok_or_else(|| RuntimeError::InvalidShape {
index,
message: format!(
"component index {i} out of bounds for len {}",
vector.len()
),
})?)
}
}
ExprNode::MatrixElement { input, row, col } => {
let (rows, cols, matrix) = self.matrix_tangent(input)?;
if row >= rows || col >= cols {
return Err(RuntimeError::InvalidShape {
index,
message: format!(
"matrix element ({row}, {col}) out of bounds for shape {rows}x{cols}"
),
});
}
Value::Scalar(matrix[row * cols + col])
}
ExprNode::MatMul { lhs, rhs } => {
let (lhs_rows, lhs_cols, lhs_value) = self.primal_matrix(lhs)?;
let (rhs_rows, rhs_cols, rhs_value) = self.primal_matrix(rhs)?;
if lhs_cols != rhs_rows {
return Err(RuntimeError::InvalidShape {
index,
message: format!(
"cannot multiply {lhs_rows}x{lhs_cols} by {rhs_rows}x{rhs_cols}"
),
});
}
let lhs_value = DMatrix::from_row_slice(lhs_rows, lhs_cols, lhs_value);
let rhs_value = DMatrix::from_row_slice(rhs_rows, rhs_cols, rhs_value);
let lhs_tangent = self.matrix_tangent_value(lhs, lhs_rows, lhs_cols)?;
let rhs_tangent = self.matrix_tangent_value(rhs, rhs_rows, rhs_cols)?;
let output = lhs_tangent * &rhs_value + lhs_value * rhs_tangent;
Value::Matrix {
rows: output.nrows(),
cols: output.ncols(),
values: matrix_values_row_major(&output),
}
}
ExprNode::MatVec { matrix, vector } => {
let (rows, cols, matrix_value) = self.primal_matrix(matrix)?;
let vector_value = self.primal_vector(vector)?;
if cols != vector_value.len() {
return Err(RuntimeError::InvalidShape {
index,
message: format!(
"cannot multiply {rows}x{cols} matrix by len {} vector",
vector_value.len()
),
});
}
let matrix_value = DMatrix::from_row_slice(rows, cols, matrix_value);
let vector_value = DVector::from_row_slice(vector_value);
let matrix_tangent = self.matrix_tangent_value(matrix, rows, cols)?;
let vector_tangent = DVector::from_vec(self.vector_tangent_value(vector, cols)?);
Value::Vector(
(matrix_tangent * vector_value + matrix_value * vector_tangent)
.iter()
.copied()
.collect(),
)
}
ExprNode::Dot { lhs, rhs } => {
let lhs_value = self.primal_vector(lhs)?;
let rhs_value = self.primal_vector(rhs)?;
if lhs_value.len() != rhs_value.len() {
return Err(RuntimeError::InvalidShape {
index,
message: format!(
"cannot dot len {} vector with len {} vector",
lhs_value.len(),
rhs_value.len()
),
});
}
let lhs_tangent = self.vector_tangent_value(lhs, lhs_value.len())?;
let rhs_tangent = self.vector_tangent_value(rhs, rhs_value.len())?;
Value::Scalar(
lhs_tangent
.iter()
.zip(rhs_value)
.map(|(lhs, rhs)| lhs * rhs)
.sum::<Complex64>()
+ lhs_value
.iter()
.zip(rhs_tangent)
.map(|(lhs, rhs)| lhs * rhs)
.sum::<Complex64>(),
)
}
ExprNode::Solve { matrix, rhs } => {
if self.cached_factors.is_some() && self.plan.cached_value_slots[index].is_none() {
self.tangents[index] = Some(Value::Vector(Vec::new()));
return Ok(());
}
let (rows, cols, matrix_value) = self.primal_matrix(matrix)?;
let solution = self.primal_vector(id)?;
let rhs_value = self.primal_vector(rhs)?;
if rows != cols || rows != rhs_value.len() {
return Err(RuntimeError::InvalidShape {
index,
message: format!(
"cannot solve {rows}x{cols} matrix against len {} vector",
rhs_value.len()
),
});
}
let matrix_tangent = self.matrix_tangent_value(matrix, rows, cols)?;
let rhs_tangent = DVector::from_vec(self.vector_tangent_value(rhs, rows)?);
let solution = DVector::from_row_slice(solution);
let tangent_rhs = rhs_tangent - matrix_tangent * solution;
let tangent = if let (Some(slot), Some((cache, row))) = (
self.plan.factor_matrix_slots[matrix.index()],
self.cached_factors,
) {
cache
.factor(slot, row)?
.solve(&tangent_rhs)
.ok_or(RuntimeError::SingularMatrix(index))?
} else if let Some(slot) = self.plan.constant_factor_slots[matrix.index()] {
self.plan.constant_factors[slot]
.get_or_init(|| DMatrix::from_row_slice(rows, cols, matrix_value).lu())
.solve(&tangent_rhs)
.ok_or(RuntimeError::SingularMatrix(index))?
} else {
let matrix_value = DMatrix::from_row_slice(rows, cols, matrix_value);
self.factors
.entry(matrix.index())
.or_insert_with(|| matrix_value.lu())
.solve(&tangent_rhs)
.ok_or(RuntimeError::SingularMatrix(index))?
};
Value::Vector(tangent.iter().copied().collect())
}
ExprNode::RealConst(_)
| ExprNode::ComplexConst(_)
| ExprNode::EventScalar(_)
| ExprNode::EventP4Component { .. } => {
return Err(RuntimeError::InvalidShape {
index,
message: "parameter-independent node appeared in a derivative lane".into(),
});
}
};
self.tangents[index] = Some(tangent);
Ok(())
}
fn primal_scalar(&self, id: ExprId) -> RuntimeResult<Complex64> {
if self.cached_factors.is_some() {
self.plan.cached_scalar_at(self.primals, id)
} else {
scalar_at(self.primals, id.index())
}
}
fn primal_vector(&self, id: ExprId) -> RuntimeResult<&[Complex64]> {
if self.cached_factors.is_some() {
self.plan.cached_vector_at(self.primals, id)
} else {
vector_at(self.primals, id.index())
}
}
fn primal_matrix(&self, id: ExprId) -> RuntimeResult<(usize, usize, &[Complex64])> {
if self.cached_factors.is_some() {
self.plan.cached_matrix_at(self.primals, id)
} else {
matrix_at(self.primals, id.index())
}
}
fn scalar_tangent(&self, id: ExprId) -> RuntimeResult<Complex64> {
match &self.tangents[id.index()] {
Some(Value::Scalar(value)) => Ok(*value),
Some(value) => Err(RuntimeError::TypeMismatch {
index: id.index(),
expected: "scalar tangent",
actual: value.kind(),
}),
None => Ok(Complex64::ZERO),
}
}
fn vector_tangent(&self, id: ExprId) -> RuntimeResult<&[Complex64]> {
match &self.tangents[id.index()] {
Some(Value::Vector(values)) => Ok(values),
Some(value) => Err(RuntimeError::TypeMismatch {
index: id.index(),
expected: "vector tangent",
actual: value.kind(),
}),
None => Err(RuntimeError::InvalidShape {
index: id.index(),
message: "inactive vector tangent requested without a target length".into(),
}),
}
}
fn vector_tangent_value(&self, id: ExprId, len: usize) -> RuntimeResult<Vec<Complex64>> {
match &self.tangents[id.index()] {
Some(Value::Vector(values)) if values.len() == len => Ok(values.clone()),
Some(Value::Vector(values)) => Err(RuntimeError::InvalidShape {
index: id.index(),
message: format!("vector tangent has len {}, expected {len}", values.len()),
}),
Some(value) => Err(RuntimeError::TypeMismatch {
index: id.index(),
expected: "vector tangent",
actual: value.kind(),
}),
None => Ok(vec![Complex64::ZERO; len]),
}
}
fn matrix_tangent(&self, id: ExprId) -> RuntimeResult<(usize, usize, &[Complex64])> {
match &self.tangents[id.index()] {
Some(Value::Matrix { rows, cols, values }) => Ok((*rows, *cols, values)),
Some(value) => Err(RuntimeError::TypeMismatch {
index: id.index(),
expected: "matrix tangent",
actual: value.kind(),
}),
None => Err(RuntimeError::InvalidShape {
index: id.index(),
message: "inactive matrix tangent requested without a target shape".into(),
}),
}
}
fn matrix_tangent_value(
&self,
id: ExprId,
rows: usize,
cols: usize,
) -> RuntimeResult<DMatrix<Complex64>> {
match &self.tangents[id.index()] {
Some(Value::Matrix {
rows: actual_rows,
cols: actual_cols,
values,
}) if *actual_rows == rows && *actual_cols == cols => {
Ok(DMatrix::from_row_slice(rows, cols, values))
}
Some(Value::Matrix {
rows: actual_rows,
cols: actual_cols,
..
}) => Err(RuntimeError::InvalidShape {
index: id.index(),
message: format!(
"matrix tangent has shape {actual_rows}x{actual_cols}, expected {rows}x{cols}"
),
}),
Some(value) => Err(RuntimeError::TypeMismatch {
index: id.index(),
expected: "matrix tangent",
actual: value.kind(),
}),
None => Ok(DMatrix::zeros(rows, cols)),
}
}
}