use crate::Real;
use itertools::izip;
use itertools::Itertools;
use nalgebra::allocator::Allocator;
use nalgebra::constraint::{DimEq, ShapeConstraint};
use nalgebra::storage::{IsContiguous, Storage, StorageMut};
use nalgebra::{
DMatrixSlice, DVector, DVectorSlice, DefaultAllocator, Dim, DimDiff, DimMin, DimMul, DimName, DimProd, DimSub,
Matrix, Matrix3, MatrixSlice, MatrixSliceMut, OMatrix, OPoint, OVector, Quaternion, Scalar, SliceStorage,
SliceStorageMut, SquareMatrix, UnitQuaternion, Vector, Vector3, U1,
};
use nalgebra_sparse::{CooMatrix, CsrMatrix};
use num::Zero;
use numeric_literals::replace_float_literals;
use std::error::Error;
use std::fmt::Display;
use std::fmt::LowerExp;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
pub use fenris_geometry::util::compute_orthonormal_vectors_3d;
pub use fenris_nested_vec::*;
use crate::allocators::{BiDimAllocator, DimAllocator};
use crate::assembly::global::CsrParAssembler;
use crate::connectivity::Connectivity;
use crate::mesh::Mesh;
use crate::nalgebra::Dynamic;
use crate::SmallDim;
pub(crate) fn clone_upper_to_lower<T, R, C, S>(matrix: &mut Matrix<T, R, C, S>)
where
T: Scalar,
R: Dim,
C: Dim,
S: StorageMut<T, R, C>,
{
for j in 0..matrix.ncols() {
for i in (j + 1)..matrix.nrows() {
matrix[(i, j)] = matrix[(j, i)].clone();
}
}
}
pub(crate) fn reshape_to_slice<T, R, C, S, R2, C2>(
matrix: &Matrix<T, R, C, S>,
shape: (R2, C2),
) -> MatrixSlice<T, R2, C2, U1, R2>
where
T: Scalar,
R: DimMul<C>,
C: Dim,
R2: DimMul<C2>,
C2: Dim,
S: Storage<T, R, C> + IsContiguous,
ShapeConstraint: DimEq<DimProd<R, C>, DimProd<R2, C2>>,
{
let (r2, c2) = shape;
assert_eq!(
matrix.nrows() * matrix.ncols(),
r2.value() * c2.value(),
"Cannot reshape with different number of elements"
);
let data_slice = matrix.as_slice();
MatrixSlice::from_slice_generic(data_slice, r2, c2)
}
pub fn coerce_col_major_slice<T, R, C, S, RSlice, CSlice>(
matrix: &Matrix<T, R, C, S>,
slice_rows: RSlice,
slice_cols: CSlice,
) -> MatrixSlice<T, RSlice, CSlice, U1, RSlice>
where
T: Scalar,
R: Dim,
RSlice: Dim,
C: Dim,
CSlice: Dim,
S: Storage<T, R, C>,
ShapeConstraint: DimEq<R, RSlice> + DimEq<C, CSlice>,
{
assert_eq!(slice_rows.value(), matrix.nrows());
assert_eq!(slice_cols.value(), matrix.ncols());
let (rstride, cstride) = matrix.strides();
assert!(
rstride == 1 && cstride == matrix.nrows(),
"Matrix must have column-major storage."
);
unsafe {
let data = SliceStorage::new_with_strides_unchecked(
&matrix.data,
(0, 0),
(slice_rows, slice_cols),
(U1::name(), slice_rows),
);
Matrix::from_data_statically_unchecked(data)
}
}
pub fn rotation_svd<T, D>(matrix: &OMatrix<T, D, D>) -> (OMatrix<T, D, D>, OVector<T, D>, OMatrix<T, D, D>)
where
T: Real,
D: DimName + DimMin<D, Output = D> + DimSub<U1>,
DefaultAllocator: Allocator<T, D>
+ Allocator<T, D, D>
+ Allocator<T, <D as DimSub<U1>>::Output>
+ Allocator<(usize, usize), D>
+ Allocator<(T, usize), D>,
{
let minus_one = T::from_f64(-1.0).unwrap();
let mut svd = matrix.clone().svd(true, true);
let min_val_idx = svd.singular_values.imin();
let mut u = svd.u.unwrap();
if u.determinant() < T::zero() {
let mut u_col = u.column_mut(min_val_idx);
u_col *= minus_one;
svd.singular_values[min_val_idx] *= minus_one;
}
let mut v_t = svd.v_t.unwrap();
if v_t.determinant() < T::zero() {
let mut v_t_row = v_t.row_mut(min_val_idx);
v_t_row *= minus_one;
svd.singular_values[min_val_idx] *= minus_one;
}
(u, svd.singular_values, v_t)
}
#[allow(non_snake_case)]
#[replace_float_literals(T::from_f64(literal).unwrap())]
pub fn apd<T: Real>(
deformation_grad: &Matrix3<T>,
initial_guess: &UnitQuaternion<T>,
max_iter: usize,
tol: T,
) -> UnitQuaternion<T> {
let F = deformation_grad;
let mut q: UnitQuaternion<T> = initial_guess.clone();
let tol_squared = tol * tol;
let mut res = T::max_value().unwrap();
let mut iter = 0;
while res > tol_squared && iter < max_iter {
let R = q.to_rotation_matrix();
let B = R.transpose() * F;
let B0 = B.column(0);
let B1 = B.column(1);
let B2 = B.column(2);
let gradient = Vector3::new(B2[1] - B1[2], B0[2] - B2[0], B1[0] - B0[1]);
let h00 = B1[1] + B2[2];
let h11 = B0[0] + B2[2];
let h22 = B0[0] + B1[1];
let h01 = (B1[0] + B0[1]) * 0.5;
let h02 = (B2[0] + B0[2]) * 0.5;
let h12 = (B2[1] + B1[2]) * 0.5;
let detH =
-(h02 * h02 * h11) + (h01 * h02 * h12) * 2.0 - (h00 * h12 * h12) - (h01 * h01 * h22) + (h00 * h11 * h22);
let factor = detH.recip() * (-0.25);
let mut omega = Vector3::zeros();
omega[0] = (h11 * h22 - h12 * h12) * gradient[0]
+ (h02 * h12 - h01 * h22) * gradient[1]
+ (h01 * h12 - h02 * h11) * gradient[2];
omega[1] = (h02 * h12 - h01 * h22) * gradient[0]
+ (h00 * h22 - h02 * h02) * gradient[1]
+ (h01 * h02 - h00 * h12) * gradient[2];
omega[2] = (h01 * h12 - h02 * h11) * gradient[0]
+ (h01 * h02 - h00 * h12) * gradient[1]
+ (h00 * h11 - h01 * h01) * gradient[2];
omega *= factor;
if detH.abs() < 1.0e-9 {
omega = -gradient;
}
let useGD = omega.dot(&gradient) > T::zero();
if useGD {
omega = &gradient * (-0.125);
}
let l_omega2 = omega.norm_squared();
let w = (1.0 - l_omega2) / (1.0 + l_omega2);
let vec = omega * (2.0 / (1.0 + l_omega2));
q = q * UnitQuaternion::new_unchecked(Quaternion::from_parts(w, vec));
iter += 1;
res = l_omega2;
}
q
}
pub fn diag_left_mul<T, D1, D2, S>(diag: &Vector<T, D1, S>, matrix: &OMatrix<T, D1, D2>) -> OMatrix<T, D1, D2>
where
T: Real,
D1: DimName,
D2: DimName,
S: Storage<T, D1>,
DefaultAllocator: Allocator<T, D1, D2>,
{
let mut result = matrix.clone();
for (i, mut row) in result.row_iter_mut().enumerate() {
row *= diag[i];
}
result
}
pub fn coerce_col_major_slice_mut<T, R, C, S, RSlice, CSlice>(
matrix: &mut Matrix<T, R, C, S>,
slice_rows: RSlice,
slice_cols: CSlice,
) -> MatrixSliceMut<T, RSlice, CSlice, U1, RSlice>
where
T: Scalar,
R: Dim,
RSlice: Dim,
C: Dim,
CSlice: Dim,
S: StorageMut<T, R, C>,
ShapeConstraint: DimEq<R, RSlice> + DimEq<C, CSlice>,
{
assert_eq!(slice_rows.value(), matrix.nrows());
assert_eq!(slice_cols.value(), matrix.ncols());
let (rstride, cstride) = matrix.strides();
assert!(
rstride == 1 && cstride == matrix.nrows(),
"Matrix must have column-major storage."
);
unsafe {
let data = SliceStorageMut::new_with_strides_unchecked(
&mut matrix.data,
(0, 0),
(slice_rows, slice_cols),
(U1::name(), slice_rows),
);
Matrix::from_data_statically_unchecked(data)
}
}
pub fn try_transmute_ref<T: 'static, U: 'static>(e: &T) -> Option<&U> {
use std::any::TypeId;
use std::mem::transmute;
if TypeId::of::<T>() == TypeId::of::<U>() {
Some(unsafe { transmute(e) })
} else {
None
}
}
pub fn try_transmute_ref_mut<T: 'static, U: 'static>(e: &mut T) -> Option<&mut U> {
use std::any::TypeId;
use std::mem::transmute;
if TypeId::of::<T>() == TypeId::of::<U>() {
Some(unsafe { transmute(e) })
} else {
None
}
}
pub fn cross_product_matrix<T: Real>(x: &Vector3<T>) -> Matrix3<T> {
Matrix3::new(T::zero(), -x[2], x[1], x[2], T::zero(), -x[0], -x[1], x[0], T::zero())
}
pub fn dump_matrix_to_file<'a, T: Scalar + Display>(
path: impl AsRef<Path>,
matrix: impl Into<DMatrixSlice<'a, T>>,
) -> Result<(), Box<dyn Error + Sync + Send>> {
let file = File::create(path.as_ref())?;
let mut writer = BufWriter::new(file);
let matrix = matrix.into();
for i in 0..matrix.nrows() {
write!(writer, "{}", matrix[(i, 0)])?;
for j in 1..matrix.ncols() {
write!(writer, " {}", matrix[(i, j)])?;
}
writeln!(writer)?;
}
writer.flush()?;
Ok(())
}
pub fn dump_mesh_connectivity_matrices<T, D, C>(
node_path: impl AsRef<Path>,
element_path: impl AsRef<Path>,
mesh: &Mesh<T, D, C>,
) -> Result<(), Box<dyn Error>>
where
T: Scalar + LowerExp,
D: DimName,
C: Sync + Connectivity,
DefaultAllocator: Allocator<T, D>,
Mesh<T, D, C>: Sync,
{
let pattern = CsrParAssembler::<usize>::default().assemble_pattern(mesh);
let nnz = pattern.nnz();
let node_matrix = CsrMatrix::try_from_pattern_and_values(pattern, vec![1.0f64; nnz])
.expect("CSR data must be valid by definition");
dump_csr_matrix_to_mm_file(node_path.as_ref(), &node_matrix).map_err(|err| err as Box<dyn Error>)?;
let mut element_node_matrix = CooMatrix::new(mesh.connectivity().len(), mesh.vertices().len());
for (i, conn) in mesh.connectivity().iter().enumerate() {
for &j in conn.vertex_indices() {
element_node_matrix.push(i, j, 1.0f64);
}
}
dump_csr_matrix_to_mm_file(element_path.as_ref(), &CsrMatrix::from(&element_node_matrix))
.map_err(|err| err as Box<dyn Error>)?;
Ok(())
}
pub fn dump_csr_matrix_to_mm_file<T: Scalar + LowerExp>(
path: impl AsRef<Path>,
matrix: &CsrMatrix<T>,
) -> Result<(), Box<dyn Error + Sync + Send>> {
let file = File::create(path.as_ref())?;
let mut writer = BufWriter::new(file);
writeln!(writer, "%%MatrixMarket matrix coordinate real general")?;
writeln!(writer, "{} {} {}", matrix.nrows(), matrix.ncols(), matrix.nnz())?;
for (i, j, v) in matrix.triplet_iter() {
writeln!(writer, "{} {} {:.e}", i + 1, j + 1, v)?;
}
writer.flush()?;
Ok(())
}
pub fn min_eigenvalue_symmetric<T, D, S>(matrix: &SquareMatrix<T, D, S>) -> T
where
T: Real,
D: Dim + DimSub<U1>,
S: Storage<T, D, D>,
DefaultAllocator:
Allocator<T, D, D> + Allocator<T, DimDiff<D, U1>> + Allocator<T, D> + Allocator<T, DimDiff<D, U1>>,
{
use std::cmp::Ordering;
matrix
.symmetric_eigenvalues()
.iter()
.min_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Less))
.unwrap()
.to_owned()
}
pub fn extract_by_node_index<T, D>(u: &[T], node_indices: &[usize]) -> DVector<T>
where
T: Scalar + Copy + Zero,
D: DimName,
{
let u = DVectorSlice::from(u);
let mut extracted = DVector::zeros(D::dim() * node_indices.len());
for (i_local, &i_global) in node_indices.iter().enumerate() {
let ui = u.rows_generic(D::dim() * i_global, D::name());
extracted
.rows_generic_mut(D::dim() * i_local, D::name())
.copy_from(&ui);
}
extracted
}
pub fn min_max_symmetric_eigenvalues<T, D, S>(matrix: &SquareMatrix<T, D, S>) -> (T, T)
where
T: Real,
D: Dim + DimSub<U1>,
S: Storage<T, D, D>,
DefaultAllocator:
Allocator<T, D, D> + Allocator<T, DimDiff<D, U1>> + Allocator<T, D> + Allocator<T, DimDiff<D, U1>>,
{
use std::cmp::Ordering;
matrix
.symmetric_eigenvalues()
.iter()
.minmax_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Less))
.into_option()
.map(|(a, b)| (*a, *b))
.unwrap()
}
pub fn condition_number_symmetric<T, D, S>(matrix: &SquareMatrix<T, D, S>) -> T
where
T: Real,
D: Dim + DimSub<U1>,
S: Storage<T, D, D>,
DefaultAllocator:
Allocator<T, D, D> + Allocator<T, DimDiff<D, U1>> + Allocator<T, D> + Allocator<T, DimDiff<D, U1>>,
{
use std::cmp::Ordering;
let (min, max) = matrix
.symmetric_eigenvalues()
.into_iter()
.cloned()
.minmax_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Less))
.into_option()
.expect("Currently don't support empty matrices");
max.abs() / min.abs()
}
#[cfg(feature = "proptest")]
pub mod proptest {
use nalgebra::{DMatrix, Point2, Scalar, Vector2};
use proptest::collection::vec;
use proptest::prelude::*;
use proptest::strategy::ValueTree;
use proptest::test_runner::{Reason, TestRunner};
pub fn point2_f64_strategy() -> impl Strategy<Value = Point2<f64>> {
vector2_f64_strategy().prop_map(|vector| Point2::from(vector))
}
pub fn vector2_f64_strategy() -> impl Strategy<Value = Vector2<f64>> {
let xrange = prop_oneof![-3.0..3.0, -100.0..100.0];
let yrange = xrange.clone();
(xrange, yrange).prop_map(|(x, y)| Vector2::new(x, y))
}
pub fn square_shape<S>(dim: S) -> impl Strategy<Value = (usize, usize)>
where
S: Strategy<Value = usize>,
{
dim.prop_map(|dim| (dim, dim))
}
#[derive(Debug, Clone, PartialEq)]
pub struct DMatrixStrategy<ElementStrategy, ShapeStrategy> {
element_strategy: ElementStrategy,
shape_strategy: ShapeStrategy,
}
impl DMatrixStrategy<(), ()> {
pub fn new() -> Self {
Self {
element_strategy: (),
shape_strategy: (),
}
}
}
impl<ElementStrategy, ShapeStrategy> DMatrixStrategy<ElementStrategy, ShapeStrategy> {
pub fn with_elements<E>(self, element_strategy: E) -> DMatrixStrategy<E, ShapeStrategy>
where
E: Strategy,
{
DMatrixStrategy {
element_strategy,
shape_strategy: self.shape_strategy,
}
}
pub fn with_shapes<S>(self, shape_strategy: S) -> DMatrixStrategy<ElementStrategy, S>
where
S: Strategy<Value = (usize, usize)>,
{
DMatrixStrategy {
element_strategy: self.element_strategy,
shape_strategy,
}
}
}
impl<ElementStrategy, ShapeStrategy> Strategy for DMatrixStrategy<ElementStrategy, ShapeStrategy>
where
ElementStrategy: Clone + 'static + Strategy,
ElementStrategy::Value: Scalar,
ShapeStrategy: Clone + 'static + Strategy<Value = (usize, usize)>,
{
type Tree = Box<dyn ValueTree<Value = Self::Value>>;
type Value = DMatrix<ElementStrategy::Value>;
fn new_tree(&self, runner: &mut TestRunner) -> Result<Self::Tree, Reason> {
let element_strategy = self.element_strategy.clone();
self.shape_strategy
.clone()
.prop_flat_map(move |(nrows, ncols)| {
let num_elements = nrows * ncols;
vec(element_strategy.clone(), num_elements)
.prop_map(move |elements| DMatrix::from_row_slice(nrows, ncols, &elements))
})
.boxed()
.new_tree(runner)
}
}
#[cfg(test)]
mod tests {
use proptest::prelude::*;
use super::DMatrixStrategy;
proptest! {
#[test]
fn dmatrix_strategy_respects_strategies(
matrix in DMatrixStrategy::new()
.with_shapes((Just(5), 2usize..=3))
.with_elements(0i32 ..= 5))
{
prop_assert_eq!(matrix.nrows(), 5);
prop_assert!(matrix.ncols() >= 2);
prop_assert!(matrix.ncols() <= 3);
prop_assert!(matrix.iter().cloned().all(|x| x >= 0 && x <= 5));
}
}
}
}
pub fn compute_interpolation<'a, 'b, T, SolutionDim>(
u: impl Into<DVectorSlice<'a, T>>,
basis: impl Into<DVectorSlice<'b, T>>,
) -> OVector<T, SolutionDim>
where
T: Real,
SolutionDim: SmallDim,
DefaultAllocator: DimAllocator<T, SolutionDim>,
{
compute_interpolation_(u.into(), basis.into())
}
fn compute_interpolation_<T, SolutionDim>(u: DVectorSlice<T>, basis: DVectorSlice<T>) -> OVector<T, SolutionDim>
where
T: Real,
SolutionDim: SmallDim,
DefaultAllocator: DimAllocator<T, SolutionDim>,
{
let s = SolutionDim::dim();
let n = basis.len();
assert_eq!(u.len(), s * n);
assert_eq!(s, SolutionDim::dim());
let u = reshape_to_slice(&u, (SolutionDim::name(), Dynamic::new(n)));
u * basis
}
pub fn compute_interpolation_gradient<'a, T, SolutionDim, GeometryDim>(
u: impl Into<DVectorSlice<'a, T>>,
basis_gradients: impl Into<DVectorSlice<'a, T>>,
) -> OMatrix<T, GeometryDim, SolutionDim>
where
T: Real,
SolutionDim: SmallDim,
GeometryDim: SmallDim,
DefaultAllocator: BiDimAllocator<T, GeometryDim, SolutionDim>,
{
compute_interpolation_gradient_(u.into(), basis_gradients.into())
}
fn compute_interpolation_gradient_<T, SolutionDim, GeometryDim>(
u: DVectorSlice<T>,
basis_gradients: DVectorSlice<T>,
) -> OMatrix<T, GeometryDim, SolutionDim>
where
T: Real,
SolutionDim: SmallDim,
GeometryDim: SmallDim,
DefaultAllocator: BiDimAllocator<T, GeometryDim, SolutionDim>,
{
let d = GeometryDim::dim();
let s = SolutionDim::dim();
let n = basis_gradients.len() / d;
assert_eq!(basis_gradients.len(), d * n);
assert_eq!(u.len(), s * n);
let u = reshape_to_slice(&u, (SolutionDim::name(), Dynamic::new(n)));
let g = reshape_to_slice(&basis_gradients, (GeometryDim::name(), Dynamic::new(n)));
let mut u_grad = OMatrix::<T, GeometryDim, SolutionDim>::zeros();
for (u_i, g_i) in izip!(u.column_iter(), g.column_iter()) {
u_grad.ger(T::one(), &g_i, &u_i, T::one());
}
u_grad
}
pub fn global_vector_from_point_fn<T, D, S, F>(points: &[OPoint<T, D>], mut f: F) -> DVector<T>
where
T: Scalar,
D: SmallDim,
S: SmallDim,
F: FnMut(&OPoint<T, D>) -> OVector<T, S>,
DefaultAllocator: DimAllocator<T, D> + DimAllocator<T, S>,
{
let mut result = Vec::with_capacity(points.len() * S::dim());
for point in points {
result.extend_from_slice(f(point).as_slice());
}
DVector::from_vec(result)
}