Mercurial > repos > non-riemannian-opt / file revision

/*!
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)]

use serde::Serialize;
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 image::{ImageFormat, ImageBuffer, Rgb};

use dist::DistToSquaredDiv2;
use fb::{forward_backward, IterInfo};
use manifold::EmbeddedManifoldPoint;
use cube::*;
use Face::*;
use zero::ZeroFn;

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

/// Program entry point
fn main() {
    simple_cube_test().unwrap()
}

/// Helper structure for saving the log into a CSV file
#[derive(Serialize)]
struct CSVLog {
    iter : usize,
    value : f64,
    face : Face,
    x : f64,
    y : f64,
    z : f64
}

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

/// A simple test on the cube
fn simple_cube_test() -> DynError {
    
    let points = [
        //OnCube::new(F1, Loc([0.5, 0.5])),
        //OnCube::new(F2, Loc([0.5, 0.5])),
        //OnCube::new(F4, Loc([0.1, 0.1])),
        OnCube::new(F1, Loc([0.5, 0.5])),
        OnCube::new(F3, Loc([0.5, 0.5])),
        OnCube::new(F2, Loc([0.5, 0.5])),
    ];

    //let x = points[0].clone();
    // OnCube::new(F3, Loc([0.5, 0.5])); goes to opposite side
    let x = OnCube::new(F3, Loc([0.5, 0.4]));
    let f = Sum::new(points.into_iter().map(DistToSquaredDiv2));
    let g = ZeroFn::new();
    let τ = 0.1;
    
    let mut logger = Logger::new();
    let logmap = |iter, IterInfo { value, point } : IterInfo<OnCube>| {
        let Loc([x,y,z]) = point.embedded_coords();
        let face = point.face();
        CSVLog { iter, value, face, x,  y, z }
    };
    let iter = AlgIteratorOptions{
        max_iter : 100,
        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̂);

    std::fs::create_dir_all(PREFIX)?;

    logger.write_csv(format!("{PREFIX}/log.txt"))?;

    for face in Face::all() {
        write_face_csv(format!("{PREFIX}/{face}"), face, 64, |x| f.apply(x) + g.apply(x))?;
        write_face_img(format!("{PREFIX}/{face}"), face, 128, |x| f.apply(x) + g.apply(x))?;
    }

    Ok(())
}

/// Writes the values of `f` on `face` of a [`OnCube`] into a PNG file
/// with resolution `n × n`.
fn write_face_img(filename : String, face : Face, n : usize, mut f : impl FnMut(&OnCube) -> f64) -> DynError {
    let mut img = ImageBuffer::new(n as u32, n as u32);
    let grid = LinSpace {
        start : Loc([0.0, 0.0]),
        end : Loc([1.0, 1.0]),
        count : [n, n]
    };
    let rawdata : Vec<_> = grid.into_iter()
                               .map(|x| f(&OnCube::new(face, x)))
                               .collect();
    let a = rawdata.iter().copied().reduce(f64::max).unwrap();
    img.pixels_mut()
        .zip(rawdata)
        .for_each(|(p, v)| {
            let t = v/a;
            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!("{filename}.png"), ImageFormat::Png)?;

    Ok(())
}

/// Writes the values of `f` on `face` of a [`OnCube`] into a CSV file
/// with resolution `n × n`.
fn write_face_csv(filename : String, face : 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(())
}