src/main.rs

/*!
Optimisation on non-Riemannian manifolds.
*/

// We use unicode. We would like to use much more of it than Rust allows.
// Live with it. Embrace it.
#![allow(uncommon_codepoints)]
#![allow(mixed_script_confusables)]
#![allow(confusable_idents)]

mod manifold;
mod fb;
mod cube;
mod cylinder;
mod dist;
mod zero;
mod scaled;
mod newton;

use serde::{Serialize, Deserialize};
use alg_tools::logger::Logger;
use alg_tools::tabledump::{TableDump, write_csv};
use alg_tools::error::DynError;
use alg_tools::lingrid::LinSpace;
use alg_tools::loc::Loc;
use alg_tools::types::*;
use alg_tools::mapping::{Sum, Apply};
use alg_tools::iterate::{AlgIteratorOptions, AlgIteratorFactory, Verbose};
use alg_tools::mapping::Mapping;
use image::{ImageFormat, ImageBuffer, Rgb};

use dist::{DistTo, DistToSquaredDiv2};
use fb::{forward_backward, IterInfo, Desc, Prox};
use manifold::{EmbeddedManifoldPoint, ManifoldPoint, FacedManifoldPoint};
use cube::OnCube;
use cylinder::{CylCoords, Cylinder, CylinderConfig, OnCylinder, normalise_angle};
#[allow(unused_imports)]
use zero::ZeroFn;
use scaled::Scaled;

/// Location for saving results
static PREFIX : &str = "res";

/// Program entry point
fn main() {
    simple_cube_test(format!("{PREFIX}/cube").as_str()).unwrap();
    simple_cylinder_test(format!("{PREFIX}/cylinder").as_str()).unwrap();
}

/// Helper structure for saving a point on a cube into a CSV file
#[derive(Serialize,Deserialize,Debug)]
struct CSVPoint<Face> {
    face : Face,
    x : f64,
    y : f64,
    z : f64
}

impl<M> From<&M> for CSVPoint<M::Face>
where M : EmbeddedManifoldPoint<EmbeddedCoords = Loc<f64, 3>> + FacedManifoldPoint {
    fn from(point : &M) -> Self {
        let Loc([x,y,z]) = point.embedded_coords();
        let face = point.face();
        CSVPoint { face, x,  y, z }
    }
}

/// Helper structure for saving a point on a cylinder into a CSV file
#[derive(Serialize,Deserialize,Debug)]
struct CSVCylPoint<Face> {
    face : Face,
    angle : f64,
    r : f64,
    z : f64
}

impl<'a> From<&'a crate::cylinder::OnCylinder<'a>> for CSVCylPoint<crate::cylinder::Face> {
    fn from(point : &'a OnCylinder<'a>) -> Self {
        let CylCoords {r, angle, z} = point.cyl_coords();
        let face = point.face();
        CSVCylPoint { face, r, angle, z }
    }
}

/// A simple test on the cube
fn simple_cube_test(dir : &str) -> DynError {
    use crate::cube::Face;
    use crate::cube::Face::*;
    
    let points = [
        OnCube::new(F1, Loc([0.5, 0.7])),
        OnCube::new(F2, Loc([0.3, 0.5])),
        OnCube::new(F4, Loc([0.9, 0.9])),
        OnCube::new(F6, Loc([0.4, 0.3])),
        OnCube::new(F4, Loc([0.3, 0.7])),
        OnCube::new(F3, Loc([0.3, 0.7])),
    ];

    let origin = OnCube::new(F4, Loc([0.5, 0.5]));

    std::fs::create_dir_all(dir)?;
    write_csv(points.iter().map(CSVPoint::from), format!("{dir}/data.csv"))?;
    write_csv(std::iter::once(&origin).map(CSVPoint::from), format!("{dir}/origin.csv"))?;

    let f = Sum::new(points.into_iter().map(DistToSquaredDiv2));
    //let g = ZeroFn::new();
    let g = Scaled::new(0.5, DistTo(origin));
    let τ = 0.1;
    
    for face in Face::all() {
        write_cube_face_csv(format!("{dir}/{face}"), face, 32, |x| f.apply(x) + g.apply(x))?;
    }
    write_cube_face_imgs(dir, 128, |x| f.apply(x) + g.apply(x))?;

    /// Helper structure for saving the log into a CSV file
    #[derive(Serialize,Deserialize,Debug)]
    struct CSVLog<Face> {
        iter : usize,
        value : f64,
        // serde is junk
        //#[serde(flatten)]
        //point : CSVPoint
        face : Face,
        x : f64,
        y : f64,
        z : f64
    }

    let logmap = |iter, IterInfo { value, point } : IterInfo<OnCube>| {
        let CSVPoint {x , y, z, face} = CSVPoint::from(&point);
        CSVLog {
            iter, value, //point : CSVPoint::from(&point)
            x, y, z, face
         }
    };

    run_and_save(dir, "x1", &f, &g, OnCube::new(F3, Loc([0.1, 0.7])), τ, logmap)?;
    run_and_save(dir, "x2", &f, &g, OnCube::new(F2, Loc([0.3, 0.1])), τ, logmap)?;
    run_and_save(dir, "x3", &f, &g, OnCube::new(F6, Loc([0.6, 0.2])), τ, logmap)
}


/// A simple test on the cube
fn simple_cylinder_test(dir : &str) -> DynError {
    let π = f64::PI;

    let cyl = Cylinder {
        radius : 1.0,
        height : 1.0,
        config : CylinderConfig { newton_iters : 100 },
        //Default::default(),
    };
    
    let points = [
        cyl.on_side(π/2.0, 0.4),
        cyl.on_side(π*3.0/2.0, -0.1),
        cyl.on_side(π, 0.2),
        cyl.on_side(-π/5.0, -0.2),
        cyl.on_side(π*2.0/3.0, -0.45),
        cyl.on_top(π*3.0/4.0, 0.7),
        cyl.on_bottom(π*5.0/4.0, 0.3),
    ];

    let origin = cyl.on_side(0.0, 0.0);

    std::fs::create_dir_all(dir)?;
    write_csv(points.iter().map(CSVCylPoint::from), format!("{dir}/data.csv"))?;
    write_csv(std::iter::once(&origin).map(CSVCylPoint::from), format!("{dir}/origin.csv"))?;

    let f = Sum::new(points.into_iter().map(DistToSquaredDiv2));
    let g = Scaled::new(4.0, DistTo(origin));
    let τ = 0.1;
    
    write_cylinder_faces_csv(dir, &cyl, 32, 32, 32, |x| f.apply(x) + g.apply(x))?;

    /// Helper structure for saving the log into a CSV file
    #[derive(Serialize,Deserialize,Debug)]
    struct CSVLog<Face> {
        iter : usize,
        value : f64,
        // serde is junk
        //#[serde(flatten)]
        //point : CSVCylPoint
        face : Face,
        angle : f64,
        r : f64,
        z : f64,
    }

    let logmap = |iter, IterInfo { value, point } : IterInfo<OnCylinder<'_>>| {
        let CSVCylPoint {r, angle, z, face} = CSVCylPoint::from(&point);
        CSVLog {
            iter, value, //point : CSVPoint::from(&point)
            r, angle, z, face
         }
    };

    run_and_save(dir, "x1", &f, &g, cyl.on_top(π*0.7, 0.8), τ, logmap)?;
    run_and_save(dir, "x2", &f, &g, cyl.on_side(-π*0.42, -0.38), τ, logmap)?;
    run_and_save(dir, "x3", &f, &g,  cyl.on_bottom(-π*0.3, 0.9), τ, logmap)

}


/// Runs [forward_backward] and saves the results.
pub fn run_and_save<F, G, M, I>(
    dir : &str,
    name : &str,
    f : &F,
    g : &G,
    x : M,
    τ : f64,
    logmap : impl Fn(usize, IterInfo<M>) -> I
) -> DynError
where M : ManifoldPoint
          + EmbeddedManifoldPoint<EmbeddedCoords = Loc<f64, 3>>
          + FacedManifoldPoint,
      F : Desc<M> +  Mapping<M, Codomain = f64>,
      G : Prox<M> +  Mapping<M, Codomain = f64>,
      I : Serialize {
    
    let mut logger = Logger::new();
    let iter = AlgIteratorOptions{
        max_iter : 20,
        verbose_iter : Verbose::Every(1),
        .. Default::default()
    }.mapped(logmap)
     .into_log(&mut logger);

    let x̂ = forward_backward(f, g, x, τ, iter);
    println!("result = {}\n{:?}", x̂.embedded_coords(), &x̂);

    logger.write_csv(format!("{dir}/{name}_log.csv"))?;

    Ok(())
}

/// Writes the values of `f` on `face` of a [`OnCube`] into a PNG file
/// with resolution `n × n`.
fn write_cube_face_imgs(
    dir : &str,
    n : usize,
    mut f : impl FnMut(&OnCube) -> f64
) -> DynError {
    use crate::cube::Face;

    let grid = LinSpace {
        start : Loc([0.0, 0.0]),
        end : Loc([1.0, 1.0]),
        count : [n, n]
    };

    let mut m = 0.0;
    let mut datas = Vec::new();

    for face in Face::all() {
        let rawdata : Vec<_> = grid.into_iter()
                                .map(|Loc([x,y])| f(&OnCube::new(face, Loc([x, 1.0-y]))))
                                .collect();
        m = rawdata.iter().copied().fold(m, f64::max);
        datas.push((face, rawdata));
    }

    for (face, rawdata) in datas {
        let mut img = ImageBuffer::new(n as u32, n as u32);
        img.pixels_mut()
            .zip(rawdata)
            .for_each(|(p, v)| {
                let t = v/m;
                // A very colourful option for bug hunting.
                //let rgb = [(50.0*t).cos(), (20.0*t).sin(), (3.0*t).cos()];
                let rgb = [1.0-t, 1.0-t, 1.0];
                *p = Rgb(rgb.map(|v| (v*(u8::RANGE_MAX as f64)) as u8))
            });

        img.save_with_format(format!("{dir}/{face}.png"), ImageFormat::Png)?;
    }

    Ok(())
}

/// Writes the values of `f` on `face` of a [`OnCube`] into a CSV file
/// with resolution `n × n`.
fn write_cube_face_csv(
    filename : String,
    face : crate::cube::Face,
    n : usize,
    mut f : impl FnMut(&OnCube) -> f64
) -> DynError {

    #[derive(Serialize)]
    struct CSVFace { u : f64, v : f64, value : f64 }

    let grid = LinSpace {
        start : Loc([0.0, 0.0]),
        end : Loc([1.0, 1.0]),
        count : [n, n]
    };

    let data = grid.into_iter()
                   .map(|p@Loc([u,v])| CSVFace{ u, v, value : f(&OnCube::new(face, p)) });

    write_csv(data, format!("{filename}.csv"))?;

    Ok(())
}

/// Writes the values of `f` on the faces of a `cylinder`.
fn write_cylinder_faces_csv<'a>(
    dir : &str,
    cyl : &'a Cylinder,
    n_height : usize,
    n_radius : usize,
    n_angle : usize,
    mut f : impl FnMut(&OnCylinder<'a>) -> f64
) -> DynError {
    let π = f64::PI;

    // Side front is [a, b] for the TikZ configuration.
    let a = -π*5.0/6.0;
    let b = π/6.0;

    #[derive(Serialize)]
    struct CSVFace { angle : f64, /*cos_angle : f64, sin_angle : f64,*/ v : f64, value : f64 }

    let mkf = |angle, v, value| CSVFace {
        angle,
        //cos_angle : angle.cos(),
        //sin_angle : angle.sin(),
        v,
        value
    };

    let side_half_grid = LinSpace {
        start : Loc([0.0, cyl.bottom_z()]),
        end : Loc([π, cyl.top_z()]),
        count : [n_angle / 2, n_height]
    };

    let side_grid = LinSpace {
        start : Loc([0.0, cyl.bottom_z()]),
        end : Loc([2.0*π, cyl.top_z()]),
        count : [n_angle, n_height]
    };
    let cap_grid = LinSpace {
        start : Loc([0.0, 0.0]),
        end : Loc([2.0*π, cyl.radius]),
        count : [n_angle, n_radius]
    };

    let side_front = side_grid.into_iter()
        .map(|Loc([angle, v])| mkf(angle, v, f(&cyl.on_side(angle, v))));
    write_csv(side_front, format!("{dir}/side.csv"))?;

    let side_front = side_half_grid.into_iter()
        .map(|Loc([angle, v])| (normalise_angle(angle + a), v))
        .map(|(angle, v)| mkf(angle, v, f(&cyl.on_side(angle, v))));
    write_csv(side_front, format!("{dir}/side_front.csv"))?;

    let side_back = side_half_grid.into_iter()
        .map(|Loc([angle, v])| (normalise_angle(angle + b), v))
        .map(|(angle, v)| mkf(angle + π, v, f(&cyl.on_side(angle + π, v))));
    write_csv(side_back, format!("{dir}/side_back.csv"))?;

    let top = cap_grid.into_iter()
        .map(|Loc([angle, v])| mkf(angle, v, f(&cyl.on_top(angle, v))));
    write_csv(top, format!("{dir}/top.csv"))?;

    let bottom = cap_grid.into_iter()
        .map(|Loc([angle, v])| mkf(angle, v, f(&cyl.on_bottom(angle, v))));
    write_csv(bottom, format!("{dir}/bottom.csv"))
    
}