rao 0.2.2

//! # rao
//!
//! `rao` - Adaptive Optics tools in Rust - is a set of fast and robust adaptive
//! optics utilities. The current scope of `rao` is for the calculation of
//! large matrices in AO, used in the configuration of real-time adaptive optics,
//! control. Specifically, we aim to provide fast, scalable, and reliable APIs for
//! generating:
//!  - `rao::IMat` - the interaction matrix between measurements and actuators,
//!  - `rao::CovMat` - the covariance matrix between measurements.
//!
//! These two matrices are typically the largest computational burden in the
//! configuration of real-time control (RTC) for AO, and also the most
//! performance-sensitive parts of the RTC.
//!
//! # Examples
//! Building an interaction matrix for a square-grid DM and a square-grid SH-WFS:
//! ```
//! use crate::rao::Matrix;
//! const N_SUBX: i32 = 8;  // 8 x 8 subapertures
//! const PITCH: f64 = 0.2;  // 0.2 metres gap between actuators
//! const COUPLING: f64 = 0.5;  // actuator cross-coupling
//!
//! // build list of measurements
//! let mut measurements = vec![];
//! for i in 0..N_SUBX {
//!     for j in 0..N_SUBX {
//!         let x0 = ((j-N_SUBX/2) as f64 + 0.5)*PITCH;
//!         let y0 = ((i-N_SUBX/2) as f64 + 0.5)*PITCH;
//!         let xz = 0.0;  // angular x-component (radians)
//!         let yz = 0.0;  // angular y-compenent (radians)
//!         // define the optical axis of subaperture
//!         let line = rao::Line::new(x0,xz,y0,yz);
//!         // slope-style measurement
//!         // x-slope
//!         measurements.push(rao::Measurement::SlopeTwoEdge{
//!             central_line: line.clone(),
//!             edge_separation: PITCH,
//!             edge_length: PITCH,
//!             npoints: 5,
//!             gradient_axis: rao::Vec2D::x_unit(),
//!             altitude: f64::INFINITY,
//!         });
//!         // y-slope
//!         measurements.push(rao::Measurement::SlopeTwoEdge{
//!             central_line: line.clone(),
//!             edge_separation: PITCH,
//!             edge_length: PITCH,
//!             npoints: 5,
//!             gradient_axis: rao::Vec2D::y_unit(),
//!             altitude: f64::INFINITY,
//!         });
//!     }
//! }
//!
//! // build list of actuators
//! let mut actuators = vec![];
//! for i in 0..(N_SUBX+1) {
//!     for j in 0..(N_SUBX+1) {
//!         let x = ((j-N_SUBX/2) as f64)*PITCH;
//!         let y = ((i-N_SUBX/2) as f64)*PITCH;
//!         actuators.push(
//!             // Gaussian influence functions
//!             rao::Actuator::Gaussian{
//!                 // std defined by coupling and pitch
//!                 sigma: rao::coupling_to_sigma(COUPLING, PITCH),
//!                 // position of actuator in 3D (z=altitude)
//!                 position: rao::Vec3D::new(x, y, 0.0),
//!                 // microns per volt of actuator
//!                 microns_per_volt: 1.0,
//!             }
//!         );
//!     }
//! }
//!
//! // instanciate imat from (actu,meas)
//! let imat = rao::IMat::new(&measurements, &actuators);
//! // serialise imat for saving
//! let data: Vec<f64> = imat.flattened_array();
//! ```

#[macro_use]
extern crate impl_ops;

pub mod utils;
pub use utils::coupling_to_sigma;
mod geometry;
pub use crate::geometry::{Line, Vec2D, Vec3D};
mod linalg;
pub use crate::linalg::Matrix;
mod core;
pub use crate::core::{CoSampleable, CovMat, IMat, Sampleable, Sampler};
use serde::{Deserialize, Serialize};

/// Common [Sampler]s in Adaptive Optics
///
/// A [Measurement] provides a scalar-valued sample of an AO system. A single
/// measurement device (e.g., a Shack Hartmann WFS) is typically comprised of
/// many [Measurement]s, e.g., `&[Measurement; N]`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum Measurement {
    /// The null measurement, always returning 0.0 regardless of the measured object.
    Zero,
    /// Phase measurement along a given [Line].
    Phase {
        /// [Line] in 3D space to trace through [Sampleable] objects.
        line: Line,
    },
    /// Slope measurement, defined by two [Line]s in 3D space.
    ///
    /// The slope measured is equal to the *sampled function at the point where
    /// `line_pos` intersects it* minus the *sampled function at the point
    /// where `line_neg` intersects it*, divided by the *distance between the
    /// two lines at zero-altitude*.
    SlopeTwoLine {
        /// positive end of slope (by convention)
        line_pos: Line,
        /// negative end of slope (by convention)
        line_neg: Line,
    },
    /// Slope measurement, defined by a [Line] in 3D space and  the sampled edges
    /// a rectangle in 2D space.
    ///
    /// Approximate the average slope over a region defined by two edges of a rectangle.
    /// The edges are defined by their length, their separation, and the axis along which
    /// they are separated (the *gradient axis*). The edges of the rectangle are sampled
    /// so that the spacing between points is uniform and the furthest distance to an
    /// unsampled point of the edge is minimal, i.e.:
    ///  - `npoints=1` => `[----x----]`,
    ///  - `npoints=2` => `[--x---x--]`,
    ///  - `npoints=3` => `[-x--x--x-]`,
    ///  - *etc*
    ///
    /// This is not a typical "linspace", which is ill-defined for 1 point, though one
    /// could implement that as a [Sampler] wrapper around [`Measurement::SlopeTwoLine`].
    SlopeTwoEdge {
        /// Principle axis of the WFS
        central_line: Line,
        /// Length of edges (e.g., subaperture height for x-slope measurement)
        edge_length: f64,
        /// Separation of edges (e.g., subaperture width for x-slope measurement)
        edge_separation: f64,
        /// Axis of gradient vector to sample, e.g., x-slope -> `Vec2D::new(1.0, 0.0)`
        gradient_axis: Vec2D,
        /// Number of points to sample along each edge (more points can be more accurate).
        npoints: u32,
        /// Altitude of target to raytrace through to.
        altitude: f64,
    },
}

/// [Measurement] is the prototypical [Sampler] type.
impl Sampler for Measurement {
    /// The [`Sampler::get_bundle`] method implementation for the [Measurement] variants
    /// should serve as an example for implementing other [Sampler] types. Inspect the
    /// source code for the reference implementation. In short:
    ///  - a [`Measurement::Zero`] variant returns an empty vector,
    ///  - a [`Measurement::Phase`] variant returns a single line with a coefficient of `1.0`,
    ///  - a [`Measurement::SlopeTwoLine`] variant returns two lines, with a positive
    ///    and negative coefficient (for the *start* and *end* of the slope), scaled by
    ///    the inverse of the ground separation of the lines, so that the resulting units
    ///    are in *sampled function units per distance unit*.
    ///  - a [`Measurement::SlopeTwoEdge`] variant returns `2 * npoints` lines, consisting of
    ///    `npoints` positive coefficients, and `npoints` negative coefficients, each scaled
    ///    by the inverse of the ground-layer separation and a factor of `1.0 / npoints as f64`,
    ///    in order to get a result which is in units of *sampled function units per distance unit*.
    fn get_bundle(&self) -> Vec<(Line, f64)> {
        match self {
            Measurement::Zero => vec![],
            Measurement::Phase { line } => vec![(line.clone(), 1.0)],
            Measurement::SlopeTwoLine { line_pos, line_neg } => {
                let coeff = 1.0 / line_pos.distance_at_ground(line_neg);
                vec![(line_pos.clone(), coeff), (line_neg.clone(), -coeff)]
            }
            Measurement::SlopeTwoEdge {
                central_line,
                edge_length,
                edge_separation,
                gradient_axis,
                npoints,
                altitude,
            } => {
                let coeff = (1.0 / f64::from(*npoints)) / edge_separation;
                let offset_vec = gradient_axis * edge_separation * 0.5;
                let point_a = edge_length * 0.5 * gradient_axis.ortho();
                let point_b = -point_a.clone();
                match *altitude {
                    f64::INFINITY => Vec2D::linspread(&point_a, &point_b, *npoints)
                        .iter()
                        .flat_map(|p| {
                            vec![
                                (central_line + (p + &offset_vec), coeff),
                                (central_line + (p - &offset_vec), -coeff),
                            ]
                        })
                        .collect(),
                    altitude => Vec2D::linspread(&point_a, &point_b, *npoints)
                        .iter()
                        .flat_map(|p| {
                            let p_alt = Vec3D::new(
                                central_line.xz * altitude,
                                central_line.yz * altitude,
                                altitude,
                            );
                            vec![
                                (
                                    Line::new_from_two_points(
                                        &((central_line + (p + &offset_vec))
                                            .position_at_altitude(0.0)
                                            + Vec3D::origin()),
                                        &p_alt,
                                    ),
                                    coeff,
                                ),
                                (
                                    Line::new_from_two_points(
                                        &((central_line + (p - &offset_vec))
                                            .position_at_altitude(0.0)
                                            + Vec3D::origin()),
                                        &p_alt,
                                    ),
                                    -coeff,
                                ),
                            ]
                        })
                        .collect(),
                }
            }
        }
    }
}

/// Common [Sampleable]s in Adaptive Optics.
///
/// An [Actuator]'s state is defined by a scalar value, so a device with `N`
/// actuatable degrees of freedom is considered as `N` different [Actuator]s,
/// e.g., `&[Actuator; N]`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum Actuator {
    /// A null actuator, making zero impact on any `Measurement`
    Zero,
    /// A circularly symmetric Gaussian actuator, centred at `position` with
    /// a specified scalar `sigma`. See [`utils::gaussian`] for more info.
    Gaussian {
        /// sigma of gaussian function in metres.
        sigma: f64,
        /// position of actuator in 3d space, z=altitude.
        position: Vec3D,
        /// microns of displacement per volt
        microns_per_volt: f64,
    },
    TipTilt {
        /// position along slope of TT actuator surface where the amplitude
        /// of the phase is equal to +1.0 units. This vector is colinear with
        /// the acuation axis. E.g., if this vector is `(1.0,0.0)`, then the
        /// response of the actuator will be a
        ///Deliberately not in arcseconds
        /// so that you have to be deliberate and careful about your units.
        unit_response: Vec2D,
    },
}

impl Sampleable for Actuator {
    fn sample(&self, line: &Line) -> f64 {
        match self {
            Self::Zero => 0.0,
            Self::Gaussian {
                sigma,
                position,
                microns_per_volt,
            } => {
                let distance = position.distance_at_altitude(line);
                utils::gaussian(distance / sigma) * microns_per_volt
            }
            Self::TipTilt { unit_response } => line.position_at_altitude(0.0).dot(unit_response),
        }
    }
}

/// Simple covariance model, this might be refactored into an enum of models.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct VonKarmanLayer {
    pub r0: f64,
    pub l0: f64,
    pub alt: f64,
    pub v: Vec2D,
}

impl VonKarmanLayer {
    /// Construct a new von Karman turbulence layer from its parameters
    #[must_use]
    pub fn new(r0: f64, l0: f64, alt: f64, v: Vec2D) -> VonKarmanLayer {
        VonKarmanLayer { r0, l0, alt, v }
    }
}

/// [`VonKarmanLayer`] is (for now) the prototypical [`CoSampleable`] object.
///
/// Perhaps confusingly, this implementation allows the cosampling of the
/// von Karman turbulence statistical model, returning the covariance between
/// two [`Line`]s intercepting that layer.
impl CoSampleable for VonKarmanLayer {
    fn cosample(&self, line_a: &Line, line_b: &Line, dt: f64) -> f64 {
        let p1 = line_a.position_at_altitude(self.alt);
        let p2 = line_b.position_at_altitude(self.alt) - dt * self.v.clone();
        utils::vk_cov((p1 - p2).norm(), self.r0, self.l0)
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Pupil {
    pub rad_outer: f64,
    pub rad_inner: f64,
    pub spider_thickness: f64,
    pub spiders: Vec<(Vec2D, Vec2D)>,
}

impl Sampleable for Pupil {
    fn sample(&self, ell: &Line) -> f64 {
        let p = ell.position_at_altitude(0.0);
        let mut out: f64 = 1.0;
        let r = p.norm();
        if r > self.rad_outer {
            out *= 0.0;
        }
        if r < self.rad_inner {
            out *= 0.0;
        }
        for spider in &self.spiders {
            if signed_distance_to_capsule(&p, &spider.0, &spider.1, self.spider_thickness / 2.0)
                < 0.0
            {
                out *= 0.0;
            }
        }
        out
    }
}

fn signed_distance_to_capsule(
    position: &Vec2D,
    capsule_start: &Vec2D,
    capsule_end: &Vec2D,
    radius: f64,
) -> f64 {
    let pa: Vec2D = position - capsule_start;
    let ba: Vec2D = capsule_end - capsule_start;
    let mut h: f64 = pa.dot(&ba) / ba.dot(&ba);
    h = h.clamp(0.0, 1.0);
    (pa - ba * h).norm() - radius
}

#[cfg(test)]
mod tests {
    use std::f64;

    use super::*;
    use approx::{assert_abs_diff_eq, assert_abs_diff_ne};

    #[test]
    fn gaussian_on_axis_phase() {
        let actuators = [Actuator::Gaussian {
            sigma: coupling_to_sigma(0.5, 1.0),
            position: Vec3D::new(0.0, 0.0, 0.0),
            microns_per_volt: 1.0,
        }];
        let measurements = [Measurement::Phase {
            line: Line::new_on_axis(0.0, 0.0),
        }];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), 1.0, epsilon = f64::EPSILON);
    }

    #[test]
    fn gaussian_off_axis_phase() {
        let actuators = [Actuator::Gaussian {
            sigma: coupling_to_sigma(0.5, 1.0),
            position: Vec3D::new(0.0, 0.0, 1000.0),
            microns_per_volt: 1.0,
        }];
        let measurements = [Measurement::Phase {
            line: Line::new(0.0, 1.0 / 1000.0, 0.0, 0.0),
        }];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), 0.5, epsilon = f64::EPSILON);
    }

    #[test]
    fn gaussian_off_axis_phase_twopoint() {
        let actuators = [Actuator::Gaussian {
            sigma: coupling_to_sigma(0.5, 1.0),
            position: Vec3D::new(0.0, 0.0, 1000.0),
            microns_per_volt: 1.0,
        }];
        let measurements = [Measurement::Phase {
            line: Line::new_from_two_points(
                &Vec3D::new(1.0, 1.0, 0.0),
                &Vec3D::new(1.0, -1.0, 2000.0),
            ),
        }];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), 0.5);
    }

    #[test]
    fn simple_symmetric() {
        let actuators = [
            Actuator::Gaussian {
                sigma: coupling_to_sigma(0.5, 1.0),
                position: Vec3D::new(0.0, 0.0, 1000.0),
                microns_per_volt: 1.0,
            },
            Actuator::Gaussian {
                sigma: coupling_to_sigma(0.5, 1.0),
                position: Vec3D::new(1.0, 0.0, 1000.0),
                microns_per_volt: 1.0,
            },
        ];
        let measurements = [
            Measurement::Phase {
                line: Line::new_on_axis(0.0, 0.0),
            },
            Measurement::Phase {
                line: Line::new(1.0, 0.0, 0.0, 0.0),
            },
        ];
        let imat = IMat::new(&measurements, &actuators);
        assert!(imat.eval(0, 0) > 0.0);
        assert_abs_diff_eq!(imat.eval(0, 0), imat.eval(1, 1));
        assert_abs_diff_eq!(imat.eval(1, 0), imat.eval(0, 1));
    }

    #[test]
    fn slope_twopoint() {
        let actuators = [Actuator::Gaussian {
            sigma: coupling_to_sigma(0.5, 1.0),
            position: Vec3D::origin(),
            microns_per_volt: 1.0,
        }];
        let line = Line::new(1.0, 0.0, 0.0, 0.0);
        let measurements = [
            Measurement::SlopeTwoLine {
                line_neg: &line + Vec2D::new(-1e-5, 0.0),
                line_pos: &line + Vec2D::new(1e-5, 0.0),
            },
            Measurement::SlopeTwoLine {
                line_neg: &line + Vec2D::new(-1e-6, 0.0),
                line_pos: &line + Vec2D::new(1e-6, 0.0),
            },
            Measurement::Phase {
                line: Line::new(1.0 + 1e-7, 0.0, 0.0, 0.0),
            },
            Measurement::Phase {
                line: Line::new(1.0 - 1e-7, 0.0, 0.0, 0.0),
            },
        ];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), imat.eval(1, 0), epsilon = 1e-8);
        assert_abs_diff_eq!(
            (imat.eval(2, 0) - imat.eval(3, 0)) / 2e-7,
            imat.eval(1, 0),
            epsilon = 1e-8
        );
    }

    #[test]
    fn slope_twoedge() {
        let actuators = [Actuator::Gaussian {
            sigma: coupling_to_sigma(0.5, 1.0),
            position: Vec3D::origin(),
            microns_per_volt: 1.0,
        }];

        let line = Line::new(1.0, 0.0, 0.0, 0.0);
        let measurements = [
            Measurement::SlopeTwoLine {
                line_neg: &line + Vec2D::new(-1e-2, 0.0),
                line_pos: &line + Vec2D::new(1e-2, 0.0),
            },
            Measurement::SlopeTwoEdge {
                central_line: line,
                edge_length: 0.0,
                edge_separation: 2e-2,
                gradient_axis: Vec2D::x_unit(),
                npoints: 100,
                altitude: f64::INFINITY,
            },
        ];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), imat.eval(1, 0), epsilon = 1e-10);
    }

    #[test]
    fn slope_tt_twopoint() {
        let actuators = [Actuator::TipTilt {
            unit_response: Vec2D::x_unit(),
        }];
        let line = Line::new(1.0, 0.0, 0.0, 0.0);
        let measurements = [
            Measurement::SlopeTwoLine {
                line_neg: &line + Vec2D::new(-1e-6, 0.0),
                line_pos: &line + Vec2D::new(1e-6, 0.0),
            },
            Measurement::SlopeTwoLine {
                line_neg: &line + Vec2D::new(-1e-7, 0.0),
                line_pos: &line + Vec2D::new(1e-7, 0.0),
            },
        ];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), imat.eval(1, 0), epsilon = 1e-8);
    }

    #[test]
    fn tt_on_axis_phase() {
        let actuators = [Actuator::TipTilt {
            unit_response: Vec2D::x_unit(),
        }];
        let measurements = [Measurement::Phase {
            line: Line::new_on_axis(0.0, 0.0),
        }];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), 0.0, epsilon = f64::EPSILON);
    }

    #[test]
    fn tt_off_axis_phase() {
        let actuators = [Actuator::TipTilt {
            unit_response: Vec2D::x_unit(),
        }];
        let measurements = [Measurement::Phase {
            line: Line::new(2.0, 0.0, 0.0, 0.0),
        }];
        let imat = IMat::new(&measurements, &actuators);
        assert_abs_diff_eq!(imat.eval(0, 0), 2.0, epsilon = f64::EPSILON);
    }

    #[test]
    fn test_vkcov() {
        let vk = VonKarmanLayer {
            r0: 0.1,
            l0: 25.0,
            alt: 1000.0,
            v: Vec2D { x: 10.0, y: 0.0 },
        };
        let line = Line::new_on_axis(0.0, 0.0);
        let a = vk.cosample(&line, &line, 0.0);
        assert_abs_diff_eq!(a, utils::vk_cov(0.0, vk.r0, vk.l0));
        assert_abs_diff_eq!(a, 856.346_613_137_351_7, epsilon = 1e-3);
        let a = vk.cosample(&line, &line, 1.0);
        assert_abs_diff_ne!(a, 856.346_613_137_351_7, epsilon = 1e-3);
    }

    #[test]
    fn test_vkcovmat() {
        let vk = VonKarmanLayer {
            r0: 0.1,
            l0: 25.0,
            alt: 1000.0,
            v: Vec2D { x: 10.0, y: 0.0 },
        };
        let measurements: Vec<Measurement> = (0..10)
            .map(|i| f64::from(i) * 0.8)
            .map(|x| Measurement::Phase {
                line: Line::new_on_axis(x, 0.0),
            })
            .collect();
        let covmat = CovMat::new(&measurements, &measurements, &vk, 1e-3);
        println!("{covmat}");
    }

    #[test]
    fn test_mixedvkcovmat() {
        let vk = VonKarmanLayer {
            r0: 0.1,
            l0: 25.0,
            alt: 1000.0,
            v: Vec2D { x: 10.0, y: 0.0 },
        };
        let line = Line::new_on_axis(0.0, 0.0);
        let measurements = [
            Measurement::Phase { line: line.clone() },
            Measurement::SlopeTwoLine {
                line_pos: &line + Vec2D::new(0.1, 0.0),
                line_neg: &line + Vec2D::new(-0.1, 0.0),
            },
            Measurement::SlopeTwoEdge {
                central_line: line.clone(),
                edge_separation: 0.2,
                edge_length: 0.2,
                gradient_axis: Vec2D::x_unit(),
                npoints: 10,
                altitude: f64::INFINITY,
            },
        ];
        let covmat = CovMat::new(&measurements, &measurements, &vk, 1e-3);
        println!("{covmat}");
    }

    #[test]
    fn make_pupil() {
        let pup = [Pupil {
            rad_outer: 4.0,
            rad_inner: 0.5,
            spider_thickness: 0.1,
            spiders: vec![
                (Vec2D::new(-4.0, -4.0), Vec2D::new(3.0, 0.0)),
                (Vec2D::new(-4.0, -4.0), Vec2D::new(3.0, 0.0)),
                (Vec2D::new(-4.0, -4.0), Vec2D::new(3.0, 0.0)),
                (Vec2D::new(-4.0, -4.0), Vec2D::new(3.0, 0.0)),
            ],
        }];
        const NPOINTS: u32 = 1000;
        let x = Vec2D::linspread(&Vec2D::new(-4.0, 0.0), &Vec2D::new(4.0, 0.0), NPOINTS);
        let y = Vec2D::linspread(&Vec2D::new(0.0, -4.0), &Vec2D::new(0.0, 4.0), NPOINTS);
        println!("{y:?}");
        let p: Vec<Measurement> = y
            .iter()
            .flat_map(|y| {
                x.iter().map(|x| Measurement::Phase {
                    line: Line::new(x.x, 0.0, y.y, 0.0),
                })
            })
            .collect();
        let pup_vec = IMat::new(&p, &pup);
        let shape = [NPOINTS as usize, NPOINTS as usize];
        let data: Vec<f64> = pup_vec.flattened_array();
        let primary_hdu = fitrs::Hdu::new(&shape, data);
        fitrs::Fits::create("/tmp/pup.fits", primary_hdu).expect("Failed to create");
    }
}