Libbum / Wafer

Parallelized 3D FDTD Schrödinger Equation Solver
MIT License
21 stars 2 forks source link

Complex numbers as option #1

Open Libbum opened 7 years ago

Libbum commented 7 years ago

Currently we are using just reals for potential and wavefunction values. Need to extend this to complex values. However, I don't want to be sending round two f64 values for everything when it's not needed. We should build in complex values only if needed from the potential.

Libbum commented 7 years ago

This is actually quite a difficult ask it seems. I've asked this question on SO. Seems that Peter may have a way forward...

Libbum commented 7 years ago

A small amount of progress on this. There are a lot of hurdles though.

extern crate num;
extern crate num_complex;
extern crate ndarray;
extern crate ndarray_parallel;

use ndarray::Array3;
use ndarray_parallel::prelude::*;
use num_complex::{Complex, Complex64};
use num::{One, Zero};
use std::ops::{Add, Mul, Sub, Div, Neg};

trait Operations<Number> : Clone +
         Add<Self, Output=Self> +
         Sub<Self, Output=Self> +
         Mul<Self, Output=Self> +
         Div<Self, Output=Self> +
         Neg<Output=Self> {
    fn array_squared(&self) -> Number;
    fn harmonic(&self, sx: usize, sy: usize, sz: usize) -> Number;
    fn harmonic_array(&mut self);
}

fn main() {
    let mut comp = Array3::<Complex64>::zeros((10,10,1));
    let mut real = Array3::<f64>::zeros((10,10,1));
    real.harmonic_array();
    comp.harmonic_array();
    println!("{:?}", real);
    println!("{:?}", comp);
}

impl Operations<f64> for Array3<f64> {
    fn array_squared(&self) -> f64 {
        self.into_par_iter().map(|&el| el * el).sum()
    }
    fn harmonic(&self, sx: usize, sy: usize, sz: usize) -> f64 {
        let numx = 50.;
        let numy = 50.;
        let numz = 50.;
        let dn = 0.01;
        let dx = (sx as f64 - numx+1.)/2.;
        let dy = (sy as f64 - numy+1.)/2.;
        let dz = (sz as f64 - numz+1.)/2.;
        let r = dn * (dx*dx+dy*dy+dz*dz);
        (r*r)/2.
    }
    fn harmonic_array(&mut self) {
        let num = self.dim();
        let dn = 0.01;
        for ((i,j,k), el) in self.indexed_iter_mut() {
            let dx = (i as f64 - num.0 as f64+1.)/2.;
            let dy = (j as f64 - num.1 as f64+1.)/2.;
            let dz = (k as f64 - num.2 as f64+1.)/2.;
            let r = dn * (dx*dx+dy*dy+dz*dz);
            *el = (r*r)/2.
        }
    }
}

impl Operations<Complex64> for Array3<Complex64> {
    fn array_squared(&self) -> Complex64 {
//        w.into_par_iter().map(|&el| el * el); //Sum not implemented for complex in iter. https://github.com/rust-num/num/issues/126
        //par iter doesn't seem to accept folds.
        self.into_iter().map(|&el| el * el).fold(Complex::new(0.0, 0.0), |acc, v| acc + v)
    }
    fn harmonic(&self, sx: usize, sy: usize, sz: usize) -> Complex64 {
        let numx = 50.;
        let numy = 50.;
        let numz = 50.;
        let dn = 0.01;
        let dx = (sx as f64 - numx+1.)/2.;
        let dy = (sy as f64 - numy+1.)/2.;
        let dz = (sz as f64 - numz+1.)/2.;
        let r = dn * (dx*dx+dy*dy+dz*dz);
        let harm = (r*r)/2.;
        Complex::new(harm, harm)
    }
    fn harmonic_array(&mut self) {
        let num = self.dim();
        let dn = 0.01;
        for ((i,j,k), el) in self.indexed_iter_mut() {
            let dx = (i as f64 - num.0 as f64+1.)/2.;
            let dy = (j as f64 - num.1 as f64+1.)/2.;
            let dz = (k as f64 - num.2 as f64+1.)/2.;
            let r = dn * (dx*dx+dy*dy+dz*dz);
            let harm = (r*r)/2.;
            *el = Complex::new(harm, harm)
        }
    }
}

So we can make traits and then run custom setups for Array3<f64> and Array3<Complex64> types, but as you can see in the example: not everything is implemented for Complex, so we may run into major performance hits. Oh well.

Libbum commented 7 years ago

Obviously the .map(|&el| el * el) can be .map(|&el| el.conj() * el) too.

Libbum commented 7 years ago

we can't have a runtime field in config choosing T ...you might need 2 function instances

This is still a major issue. I can't figure it out.

Libbum commented 7 years ago

So. Seems like bluss has been working on this issue already. It's not solved, nor stable, but it works for my purposes I think.

complexfloat. May want to fork it and edit it directly if needed.

extern crate complexfloat;
extern crate num_complex;
extern crate ndarray;

use ndarray::Array3;
use num_complex::Complex;
use complexfloat::ComplexFloat;

fn output<F: ComplexFloat>(x: F) {
    println!("{:?}", x);
    println!("{:.2e} + i{:.2e}", x.real(), x.imag());
    println!("{}", std::mem::size_of_val(&x));
}

fn main() {
    let f = 3.14159;
    let z = Complex::new(1., 2.);
    output(f);
    output(z);

    let comp = Array3::<Complex<f64>>::from_elem((10,10,1), Complex::<f64>::new(1., 2.));
    let real = Array3::<f64>::from_elem((10,10,1), 2.);

    let comp_sq = array_squared(&comp);
    let real_sq = array_squared(&real);

    println!("{}, ({})", real_sq, std::mem::size_of_val(&real_sq));
    println!("{}, ({})", comp_sq, std::mem::size_of_val(&comp_sq));
}

fn array_squared<F: ComplexFloat>(w: &Array3<F>) -> F {
    w.into_iter().map(|&el| el.conj() * el).fold(F::zero(), |acc, v| acc + v)
}
Libbum commented 7 years ago

Num is still not ready for this to happen. I'm helping update the crate as much as I can, but for the moment this needs to be pushed to the v0.2 milestone. There's also a possible proof that Jared has read that we may never actually need to include complex potentials. Will hunt for that too.