1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

//! # Analytic diffraction image generator
//!
//! Functions that build a diffraction image of a helical structure using an analytic solution

use crate::bessel_utils::{ Jn };
use crate::helix::Helix;

use ndarray::{Array2, Array1, Array, ArrayView};
use std::f64::consts::PI;
use num:: {Complex, complex::Complex64, traits::Pow };


/// Computes Analytic diffraction pattern for helix object
///
/// returns a vector, with values in order (R,G,B,A) and then next pixel etc.
/// ```
/// use helixiser::helix::{ Helix, Handedness };
/// use helixiser::diffraction_analytic::diff_analytic;
///
/// let strand_1 = Helix {
///         radius: 1.,
///         rise: 0.34,
///         frequency: 10.,
///         unit_size: 0.18,
///         offset: 0.,
///         rotation: 0.,
///         handedness: Handedness::Right,
/// };
///
/// let strand_2 = Helix {
///        rotation: 143.,
///         ..strand_1  // copy remaining fields over from strand 1
/// };
/// //image as RGBA vector will be 512x512x4 (height x width x RGBA)elements
/// let image_as_RGBA_vector: Vec<f64> = diff_analytic(vec![strand_1, strand_2], 6, 2, 0.01, 512);
/// ```
pub fn diff_analytic(helices: Vec<Helix>, n_range: u8, m_range: u8, scale: f64, raster_size: u32) -> Vec<f64> {
    let num_pix: usize = (raster_size * raster_size) as usize;

    let mut m_vals: Vec<f64> = Vec::new();
    for i in 0..=m_range {
        if i == 0 { m_vals.push(i as f64) } else {
            // add +m and -m
            m_vals.push(i as f64);
            m_vals.push(-(i as f64));
        }
    };

    // intialise the 2D array that represents our image.
    let mut image: Array2<Complex64> = Array::zeros((raster_size as usize, raster_size as usize));

    // get a set of coordinates [-rastersize/2 to rastersize/2]
    let coords: Array1<f64> = Array::range(0., raster_size as f64, 1 as f64) - (raster_size / 2) as f64;

    let mut z_line: f64;  // the layerline coordinate
    let mut z_draw_coord: f64;  // rounded value of z_line (for drawing in image)
    let mut bessel: Array1<f64> = Array::zeros(raster_size as usize);
    let mut Ufac: Array1<Complex64> = Array::zeros(raster_size as usize);

    let mut phi_j: f64;  // phase change due to angular offset helix
    let mut phi_z: f64;  // phase change due to offset along helix axis (z)

    //loop over different helices
    for helix in &helices {
        let pitch: f64 = helix.pitch();
        let midpoint: f64 = raster_size as f64 / 2.;
        phi_j = helix.rotation_to_rad(); // rotation due to rotation

        // compute analytic FFT solution and modify <image>
        for n in 0..=n_range {
            // we compute the bessel function (depends only on n, not m)
            for index in 0..raster_size as usize {
                bessel[index] = Jn(n.into(), 2. * PI * coords[index].abs() * helix.radius * scale);
            }

            for m in &m_vals {
                // we determine the layerline position
                z_line = (n as f64) / (pitch * scale) + m / (helix.rise * scale);
                z_draw_coord = z_line.round();  // get the rounded value, for plotting

                phi_z = 2. * PI * z_line * helix.offset * scale;
                for index in 0..raster_size as usize {
                    Ufac[index] = bessel[index] * (Complex::new(0.0, (n as f64) * (PI / 2. - phi_j) + phi_z)).exp();
                }

                //web_sys::console::log_1(&format!("N is {n}, M is {m}, and z_line is {z}", n=n, m=m, z=z_line).into());

                // check if we are still within the limits of the rasterised image
                // if so, add the layerline amplitudes (Ufac)
                if z_draw_coord > -(raster_size as f64) / 2. && z_draw_coord < (raster_size as f64) / 2. {
                    let mut line = image.slice_mut(s![(midpoint + z_draw_coord) as usize, ..]);
                    line += &ArrayView::from(&Ufac);
                    if n != 0 {
                        let mut line = image.slice_mut(s![(midpoint - z_draw_coord) as usize, ..]);
                        line += &ArrayView::from(&Ufac);
                    }
                }
            }
        }
    }

    // flatten the array
    let image_flat = image.into_shape((num_pix, 1)).unwrap();
    let normalisation = helices.len() as f64;  // ensure that more helices does not mean higher intensities

    let intensity = image_flat.mapv(
        //  we compute the intensity, which is |amplitude|^2.
        //  We clip some values (>256), but it's probably best for now with limited dynamic range
        |amplitude| (amplitude / normalisation).norm().pow(2) * 511.
    );

    let mut out: Vec<f64> = vec![0 as f64; 4 * num_pix];
    for i in 0..num_pix {
        out[i * 4] = intensity[[i, 0]]; //image_flat[[i, 0]]; // set Red
        out[i * 4 + 1] = intensity[[i, 0]]; // set Green
        out[i * 4 + 2] = intensity[[i, 0]]; // set Blue
        out[i * 4 + 3] = 255 as f64;  // set Alpha
    }
    out
}