<commit_before>extern crate neat;
extern crate rand;
extern crate graph_neighbor_matching;
extern crate graph_io_gml as gml;

use neat::population::{Population, Rated, Unrated, Individual, Runner};
use neat::network::{NetworkGenome, NetworkGenomeDistance, LinkGeneListDistance, LinkGene, NodeGene, NodeType};
use neat::innovation::{Innovation, InnovationContainer};
use neat::fitness::Fitness;
use neat::traits::Mate;
use neat::crossover::ProbabilisticCrossover;
use neat::prob::is_probable;
use graph_neighbor_matching::similarity_max_degree;
use graph_neighbor_matching::graph::{OwnedGraph, GraphBuilder};
use std::collections::BTreeMap;
use rand::{Rng, Closed01};
use rand::distributions::{WeightedChoice, Weighted, IndependentSample};
use std::marker::PhantomData;

fn load_graph(graph_file: &str) -> OwnedGraph {
    use std::fs::File;
    use std::io::Read;

    let graph_s = {
        let mut graph_file = File::open(graph_file).unwrap();
        let mut graph_s = String::new();
        let _ = graph_file.read_to_string(&mut graph_s).unwrap();
        graph_s
    };

    let graph = gml::parse_gml(&graph_s,
                               &|_| -> Option<()> { Some(()) },
                               &|_| -> Option<()> { Some(()) })
                    .unwrap();
    OwnedGraph::from_petgraph(&graph)
}

fn genome_to_graph(genome: &NetworkGenome) -> OwnedGraph {
    let mut builder = GraphBuilder::new();

    for (&innov, node) in genome.node_genes.map.iter() {
        // make sure the node exists, even if there are no connection to it.
        let _ = builder.add_or_replace_node(innov.get());
    }

    for link in genome.link_genes.map.values() {
        if link.active {
            builder.add_edge_unweighted(link.source_node_gene.get(), link.target_node_gene.get());
        }
    }

    return builder.graph();
}

#[derive(Debug)]
struct FitnessEvaluator {
    target_graph: OwnedGraph,
}

impl FitnessEvaluator {
    // A larger fitness means "better"
    fn fitness(&self, genome: &NetworkGenome) -> f32 {
        let graph = genome_to_graph(genome);
        similarity_max_degree(&graph, &self.target_graph, 50, 0.01).get()
    }
}

struct Environment {
    node_innovation_counter: Innovation,
    link_innovation_counter: Innovation,
    // (src_node, target_node) -> link_innovation
    link_innovation_cache: BTreeMap<(Innovation, Innovation), Innovation>,
}

impl Environment {
    fn new() -> Environment {
        Environment {
            node_innovation_counter: Innovation::new(0),
            link_innovation_counter: Innovation::new(0),
            link_innovation_cache: BTreeMap::new(),
        }
    }

    fn random_link_weight<R: Rng>(&mut self, rng: &mut R) -> f64 {
        // XXX Choose a weight between -1 and 1?
        rng.gen()
    }

    fn insert_new_link<R: Rng>(&mut self, genome: &mut NetworkGenome, source_node: Innovation, target_node: Innovation, rng: &mut R) {
        let link_gene = LinkGene {
            source_node_gene: source_node,
            target_node_gene: target_node,
            weight: self.random_link_weight(rng),
            active: true,
        };
        self.insert_new_link_gene(genome, link_gene);
    }

    fn insert_new_link_gene(&mut self, genome: &mut NetworkGenome, link_gene: LinkGene) {
        let new_link_innovation = self.new_link_innovation(link_gene.source_node_gene, link_gene.target_node_gene);
        genome.link_genes.insert(new_link_innovation, link_gene);
    }

    fn mutate_add_connection<R: Rng>(&mut self, genome: &NetworkGenome, rng: &mut R) -> Option<NetworkGenome> {
        genome.find_unconnected_pair(rng).map(|(src, target)| {
            let mut offspring = genome.clone();
            // Add new link to the offspring genome
            self.insert_new_link(&mut offspring, src, target, rng);
            offspring
        })
    }

    /// choose a random link. split it in half.
    fn mutate_add_node<R: Rng>(&mut self, genome: &NetworkGenome, rng: &mut R) -> Option<NetworkGenome> {
        if let Some(link_innov) = genome.find_random_active_link_gene(rng) {
            // split link in half.
            let mut offspring = genome.clone();
            let new_node_innovation = self.new_node_innovation();
            // add new node
            offspring.node_genes.insert(new_node_innovation, NodeGene{node_type: NodeType::Hidden});
            // disable `link_innov` in offspring
            // we keep this gene (but disable it), because this allows us to have a structurally
            // compatible genome to the old one, as disabled genes are taken into account for
            // the genomic distance measure.
            offspring.link_genes.get_mut(&link_innov).unwrap().disable();
            // add two new link innovations with the new node in the middle.
            // XXX: Choose random weights? Or split weight? We use random weights for now.
            let (orig_src_node, orig_target_node) = {
                let orig_link = offspring.link_genes.get(&link_innov).unwrap();
                (orig_link.source_node_gene, orig_link.target_node_gene)
            };
            self.insert_new_link(&mut offspring, orig_src_node, new_node_innovation, rng);
            self.insert_new_link(&mut offspring, new_node_innovation, orig_target_node, rng);
            Some(offspring)
        } else {
            None
        }
    }

    fn new_node_innovation(&mut self) -> Innovation {
        self.node_innovation_counter.next().unwrap()
    }

    fn new_link_innovation(&mut self, source_node_gene: Innovation, target_node_gene: Innovation) -> Innovation {
        let key = (source_node_gene, target_node_gene);
        if let Some(&cached_innovation) = self.link_innovation_cache.get(&key) {
            return cached_innovation;
        }
        let new_innovation = self.link_innovation_counter.next().unwrap();
        self.link_innovation_cache.insert(key, new_innovation);
        new_innovation
    }

    // Generates a NetworkGenome with `n_inputs` input nodes and `n_outputs` output nodes.
    // The genome will not have any link nodes.
    fn generate_genome(&mut self, n_inputs: usize, n_outputs: usize) -> NetworkGenome {
        assert!(n_inputs > 0 && n_outputs > 0);
        let mut nodes = InnovationContainer::new();
        for _ in 0..n_inputs {
            nodes.insert(self.node_innovation_counter.next().unwrap(),
                         NodeGene { node_type: NodeType::Input });
        }
        assert!(nodes.len() == n_inputs);
        for _ in 0..n_outputs {
            nodes.insert(self.node_innovation_counter.next().unwrap(),
                         NodeGene { node_type: NodeType::Output });
        }
        assert!(nodes.len() == n_inputs + n_outputs);
        NetworkGenome {
            link_genes: InnovationContainer::new(),
            node_genes: nodes,
        }
    }
}

const POP_SIZE: usize = 100;
const INPUTS: usize = 2;
const OUTPUTS: usize = 3;

#[derive(Debug, Copy, Clone)]
enum RecombinationMethod {
    MutateWeight,
    MutateStructure,
    Crossover,
}

struct Mater<'a> {
    w_mutate_weight: u32,
    w_mutate_structure: u32,
    w_crossover: u32,
    env: &'a mut Environment,
}

impl<'a> Mate<NetworkGenome> for Mater<'a> {
    // Add an argument that descibes whether both genomes are of equal fitness.
    // Pass individual, which includes the fitness.
    fn mate<R: Rng>(&mut self,
                    parent_left: &NetworkGenome,
                    parent_right: &NetworkGenome,
                    rng: &mut R)
                    -> NetworkGenome {
        assert!(self.w_mutate_weight + self.w_mutate_structure + self.w_crossover  > 0);
        let mut items = [
            Weighted{weight: self.w_mutate_weight, item: RecombinationMethod::MutateWeight},
            Weighted{weight: self.w_mutate_structure, item: RecombinationMethod::MutateStructure},
            Weighted{weight: self.w_crossover, item: RecombinationMethod::Crossover},
        ];
        let wc = WeightedChoice::new(&mut items);

        match wc.ind_sample(rng) {
            RecombinationMethod::MutateWeight => {
                // TODO
               parent_left.clone()
            }
            RecombinationMethod::MutateStructure => {
                if is_probable(&Closed01(0.5), rng) {
                    // AddConnection
                    let offspring = self.env.mutate_add_connection(parent_left, rng).
                        or_else(|| self.env.mutate_add_connection(parent_right, rng)).
                        unwrap_or_else(|| parent_left.clone());
                    offspring
                } else {
                    // AddNode (split existing connection)
                    let offspring = self.env.mutate_add_node(parent_left, rng).
                        or_else(|| self.env.mutate_add_node(parent_right, rng)).
                        unwrap_or_else(|| parent_left.clone());
                    offspring
                }
            }
            RecombinationMethod::Crossover => {
                let x = ProbabilisticCrossover {
                    prob_match_left: Closed01(0.5), // NEAT always selects a random parent for matching genes
                    prob_disjoint_left: Closed01(0.9),
                    prob_excess_left: Closed01(0.9),
                    prob_disjoint_right: Closed01(0.15),
                    prob_excess_right: Closed01(0.15),
                };

                NetworkGenome::crossover(parent_left, parent_right, &x, rng)
            }
        }

    }
}

fn main() {
    let mut rng = rand::thread_rng();

    let fitness_evaluator = FitnessEvaluator { target_graph: load_graph("examples/jeffress.gml") };
    println!("{:?}", fitness_evaluator);

    // start with minimal random topology.
    let mut env = Environment::new();

    let template_genome = env.generate_genome(INPUTS, OUTPUTS);

    println!("{:#?}", template_genome);

    let mut initial_pop = Population::<_, Unrated>::new();

    for _ in 0..POP_SIZE {
        // Add a single link gene! This is required, otherwise we can't determine
        // correctly an InnovationRange.
        let genome = env.mutate_add_connection(&template_genome, &mut rng).unwrap();

        initial_pop.add_genome(Box::new(genome));
    }
    assert!(initial_pop.len() == POP_SIZE);

    let mut mater = Mater {
        w_mutate_weight: 0,
        w_mutate_structure: 30,
        w_crossover: 70,
        env: &mut env,
    };

    let compatibility = NetworkGenomeDistance {
        l: LinkGeneListDistance {
            excess: 1.0,
            disjoint: 1.0,
            weight: 0.0,
        },
    };

    let mut runner = Runner {
        pop_size: POP_SIZE,
        elite_percentage: Closed01(0.05),
        selection_percentage: Closed01(0.2),
        tournament_k: 3,
        compatibility_threshold: 1.0,
        compatibility: &compatibility,
        mate: &mut mater,
        fitness: &|genome| Fitness::new(fitness_evaluator.fitness(genome) as f64),
        _marker: PhantomData,
    };

    let (iter, new_pop) = runner.run(initial_pop, &|iter, pop| {
        iter >= 100 || pop.max_fitness().unwrap().get() > 0.99
    }, &mut rng);

    let new_pop = new_pop.sort();

    println!("iter: {}", iter);
    println!("{:#?}", new_pop);
}
<commit_msg>Explicit use graph-neighbor-matching calls. We will extend it in a later commit<commit_after>extern crate neat;
extern crate rand;
extern crate graph_neighbor_matching;
extern crate graph_io_gml as gml;

use neat::population::{Population, Rated, Unrated, Individual, Runner};
use neat::network::{NetworkGenome, NetworkGenomeDistance, LinkGeneListDistance, LinkGene, NodeGene, NodeType};
use neat::innovation::{Innovation, InnovationContainer};
use neat::fitness::Fitness;
use neat::traits::Mate;
use neat::crossover::ProbabilisticCrossover;
use neat::prob::is_probable;
use graph_neighbor_matching::{SimilarityMatrix, ScoreNorm, IgnoreNodeColors};
use graph_neighbor_matching::graph::{OwnedGraph, GraphBuilder};
use std::collections::BTreeMap;
use rand::{Rng, Closed01};
use rand::distributions::{WeightedChoice, Weighted, IndependentSample};
use std::marker::PhantomData;

fn load_graph(graph_file: &str) -> OwnedGraph {
    use std::fs::File;
    use std::io::Read;

    let graph_s = {
        let mut graph_file = File::open(graph_file).unwrap();
        let mut graph_s = String::new();
        let _ = graph_file.read_to_string(&mut graph_s).unwrap();
        graph_s
    };

    let graph = gml::parse_gml(&graph_s,
                               &|_| -> Option<()> { Some(()) },
                               &|_| -> Option<()> { Some(()) })
                    .unwrap();
    OwnedGraph::from_petgraph(&graph)
}

fn genome_to_graph(genome: &NetworkGenome) -> OwnedGraph {
    let mut builder = GraphBuilder::new();

    for (&innov, node) in genome.node_genes.map.iter() {
        // make sure the node exists, even if there are no connection to it.
        let _ = builder.add_or_replace_node(innov.get());
    }

    for link in genome.link_genes.map.values() {
        if link.active {
            builder.add_edge_unweighted(link.source_node_gene.get(), link.target_node_gene.get());
        }
    }

    return builder.graph();
}

#[derive(Debug)]
struct FitnessEvaluator {
    target_graph: OwnedGraph,
}

impl FitnessEvaluator {
    // A larger fitness means "better"
    fn fitness(&self, genome: &NetworkGenome) -> f32 {
        let graph = genome_to_graph(genome);
        let mut s = SimilarityMatrix::new(&graph, &self.target_graph, IgnoreNodeColors);
        s.iterate(50, 0.01);
        s.score_optimal_sum_norm(None, ScoreNorm::MaxDegree).get()
    }
}

struct Environment {
    node_innovation_counter: Innovation,
    link_innovation_counter: Innovation,
    // (src_node, target_node) -> link_innovation
    link_innovation_cache: BTreeMap<(Innovation, Innovation), Innovation>,
}

impl Environment {
    fn new() -> Environment {
        Environment {
            node_innovation_counter: Innovation::new(0),
            link_innovation_counter: Innovation::new(0),
            link_innovation_cache: BTreeMap::new(),
        }
    }

    fn random_link_weight<R: Rng>(&mut self, rng: &mut R) -> f64 {
        // XXX Choose a weight between -1 and 1?
        rng.gen()
    }

    fn insert_new_link<R: Rng>(&mut self, genome: &mut NetworkGenome, source_node: Innovation, target_node: Innovation, rng: &mut R) {
        let link_gene = LinkGene {
            source_node_gene: source_node,
            target_node_gene: target_node,
            weight: self.random_link_weight(rng),
            active: true,
        };
        self.insert_new_link_gene(genome, link_gene);
    }

    fn insert_new_link_gene(&mut self, genome: &mut NetworkGenome, link_gene: LinkGene) {
        let new_link_innovation = self.new_link_innovation(link_gene.source_node_gene, link_gene.target_node_gene);
        genome.link_genes.insert(new_link_innovation, link_gene);
    }

    fn mutate_add_connection<R: Rng>(&mut self, genome: &NetworkGenome, rng: &mut R) -> Option<NetworkGenome> {
        genome.find_unconnected_pair(rng).map(|(src, target)| {
            let mut offspring = genome.clone();
            // Add new link to the offspring genome
            self.insert_new_link(&mut offspring, src, target, rng);
            offspring
        })
    }

    /// choose a random link. split it in half.
    fn mutate_add_node<R: Rng>(&mut self, genome: &NetworkGenome, rng: &mut R) -> Option<NetworkGenome> {
        if let Some(link_innov) = genome.find_random_active_link_gene(rng) {
            // split link in half.
            let mut offspring = genome.clone();
            let new_node_innovation = self.new_node_innovation();
            // add new node
            offspring.node_genes.insert(new_node_innovation, NodeGene{node_type: NodeType::Hidden});
            // disable `link_innov` in offspring
            // we keep this gene (but disable it), because this allows us to have a structurally
            // compatible genome to the old one, as disabled genes are taken into account for
            // the genomic distance measure.
            offspring.link_genes.get_mut(&link_innov).unwrap().disable();
            // add two new link innovations with the new node in the middle.
            // XXX: Choose random weights? Or split weight? We use random weights for now.
            let (orig_src_node, orig_target_node) = {
                let orig_link = offspring.link_genes.get(&link_innov).unwrap();
                (orig_link.source_node_gene, orig_link.target_node_gene)
            };
            self.insert_new_link(&mut offspring, orig_src_node, new_node_innovation, rng);
            self.insert_new_link(&mut offspring, new_node_innovation, orig_target_node, rng);
            Some(offspring)
        } else {
            None
        }
    }

    fn new_node_innovation(&mut self) -> Innovation {
        self.node_innovation_counter.next().unwrap()
    }

    fn new_link_innovation(&mut self, source_node_gene: Innovation, target_node_gene: Innovation) -> Innovation {
        let key = (source_node_gene, target_node_gene);
        if let Some(&cached_innovation) = self.link_innovation_cache.get(&key) {
            return cached_innovation;
        }
        let new_innovation = self.link_innovation_counter.next().unwrap();
        self.link_innovation_cache.insert(key, new_innovation);
        new_innovation
    }

    // Generates a NetworkGenome with `n_inputs` input nodes and `n_outputs` output nodes.
    // The genome will not have any link nodes.
    fn generate_genome(&mut self, n_inputs: usize, n_outputs: usize) -> NetworkGenome {
        assert!(n_inputs > 0 && n_outputs > 0);
        let mut nodes = InnovationContainer::new();
        for _ in 0..n_inputs {
            nodes.insert(self.node_innovation_counter.next().unwrap(),
                         NodeGene { node_type: NodeType::Input });
        }
        assert!(nodes.len() == n_inputs);
        for _ in 0..n_outputs {
            nodes.insert(self.node_innovation_counter.next().unwrap(),
                         NodeGene { node_type: NodeType::Output });
        }
        assert!(nodes.len() == n_inputs + n_outputs);
        NetworkGenome {
            link_genes: InnovationContainer::new(),
            node_genes: nodes,
        }
    }
}

const POP_SIZE: usize = 100;
const INPUTS: usize = 2;
const OUTPUTS: usize = 3;

#[derive(Debug, Copy, Clone)]
enum RecombinationMethod {
    MutateWeight,
    MutateStructure,
    Crossover,
}

struct Mater<'a> {
    w_mutate_weight: u32,
    w_mutate_structure: u32,
    w_crossover: u32,
    env: &'a mut Environment,
}

impl<'a> Mate<NetworkGenome> for Mater<'a> {
    // Add an argument that descibes whether both genomes are of equal fitness.
    // Pass individual, which includes the fitness.
    fn mate<R: Rng>(&mut self,
                    parent_left: &NetworkGenome,
                    parent_right: &NetworkGenome,
                    rng: &mut R)
                    -> NetworkGenome {
        assert!(self.w_mutate_weight + self.w_mutate_structure + self.w_crossover  > 0);
        let mut items = [
            Weighted{weight: self.w_mutate_weight, item: RecombinationMethod::MutateWeight},
            Weighted{weight: self.w_mutate_structure, item: RecombinationMethod::MutateStructure},
            Weighted{weight: self.w_crossover, item: RecombinationMethod::Crossover},
        ];
        let wc = WeightedChoice::new(&mut items);

        match wc.ind_sample(rng) {
            RecombinationMethod::MutateWeight => {
                // TODO
               parent_left.clone()
            }
            RecombinationMethod::MutateStructure => {
                if is_probable(&Closed01(0.5), rng) {
                    // AddConnection
                    let offspring = self.env.mutate_add_connection(parent_left, rng).
                        or_else(|| self.env.mutate_add_connection(parent_right, rng)).
                        unwrap_or_else(|| parent_left.clone());
                    offspring
                } else {
                    // AddNode (split existing connection)
                    let offspring = self.env.mutate_add_node(parent_left, rng).
                        or_else(|| self.env.mutate_add_node(parent_right, rng)).
                        unwrap_or_else(|| parent_left.clone());
                    offspring
                }
            }
            RecombinationMethod::Crossover => {
                let x = ProbabilisticCrossover {
                    prob_match_left: Closed01(0.5), // NEAT always selects a random parent for matching genes
                    prob_disjoint_left: Closed01(0.9),
                    prob_excess_left: Closed01(0.9),
                    prob_disjoint_right: Closed01(0.15),
                    prob_excess_right: Closed01(0.15),
                };

                NetworkGenome::crossover(parent_left, parent_right, &x, rng)
            }
        }

    }
}

fn main() {
    let mut rng = rand::thread_rng();

    let fitness_evaluator = FitnessEvaluator { target_graph: load_graph("examples/jeffress.gml") };
    println!("{:?}", fitness_evaluator);

    // start with minimal random topology.
    let mut env = Environment::new();

    let template_genome = env.generate_genome(INPUTS, OUTPUTS);

    println!("{:#?}", template_genome);

    let mut initial_pop = Population::<_, Unrated>::new();

    for _ in 0..POP_SIZE {
        // Add a single link gene! This is required, otherwise we can't determine
        // correctly an InnovationRange.
        let genome = env.mutate_add_connection(&template_genome, &mut rng).unwrap();

        initial_pop.add_genome(Box::new(genome));
    }
    assert!(initial_pop.len() == POP_SIZE);

    let mut mater = Mater {
        w_mutate_weight: 0,
        w_mutate_structure: 30,
        w_crossover: 70,
        env: &mut env,
    };

    let compatibility = NetworkGenomeDistance {
        l: LinkGeneListDistance {
            excess: 1.0,
            disjoint: 1.0,
            weight: 0.0,
        },
    };

    let mut runner = Runner {
        pop_size: POP_SIZE,
        elite_percentage: Closed01(0.05),
        selection_percentage: Closed01(0.2),
        tournament_k: 3,
        compatibility_threshold: 1.0,
        compatibility: &compatibility,
        mate: &mut mater,
        fitness: &|genome| Fitness::new(fitness_evaluator.fitness(genome) as f64),
        _marker: PhantomData,
    };

    let (iter, new_pop) = runner.run(initial_pop, &|iter, pop| {
        iter >= 100 || pop.max_fitness().unwrap().get() > 0.99
    }, &mut rng);

    let new_pop = new_pop.sort();

    println!("iter: {}", iter);
    println!("{:#?}", new_pop);
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate ipc_channel;
extern crate test;

/// Allows doing multiple inner iterations per bench.iter() run.
///
/// This is mostly to amortise the overhead of spawning a thread in the benchmark
/// when sending larger messages (that might be fragmented).
///
/// Note that you need to compensate the displayed results
/// for the proportionally longer runs yourself,
/// as the benchmark framework doesn't know about the inner iterations...
const ITERATIONS: usize = 1;

mod platform {
    extern crate crossbeam;

    use ipc_channel::platform;
    use ITERATIONS;

    use std::sync::{mpsc, Mutex};
    use test;

    #[bench]
    fn create_channel(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..ITERATIONS {
                platform::channel().unwrap();
            }
        });
    }

    fn bench_transfer_data(b: &mut test::Bencher, size: usize) {
        let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
        let (tx, rx) = platform::channel().unwrap();

        let (wait_tx, wait_rx) = mpsc::channel();
        let wait_rx = Mutex::new(wait_rx);

        if size > platform::OsIpcSender::get_max_fragment_size() {
            b.iter(|| {
                crossbeam::scope(|scope| {
                    let tx = tx.clone();
                    scope.spawn(|| {
                        let wait_rx = wait_rx.lock().unwrap();
                        let tx = tx;
                        for _ in 0..ITERATIONS {
                            tx.send(&data, vec![], vec![]).unwrap();
                            if ITERATIONS > 1 {
                                // Prevent beginning of the next send
                                // from overlapping with receive of last fragment,
                                // as otherwise results of runs with a large tail fragment
                                // are significantly skewed.
                                wait_rx.recv().unwrap();
                            }
                        }
                    });
                    for _ in 0..ITERATIONS {
                        rx.recv().unwrap();
                        if ITERATIONS > 1 {
                            wait_tx.send(()).unwrap();
                        }
                    }
                    // For reasons mysterious to me,
                    // not returning a value *from every branch*
                    // adds some 100 ns or so of overhead to all results --
                    // which is quite significant for very short tests...
                    0
                })
            });
        } else {
            b.iter(|| {
                for _ in 0..ITERATIONS {
                    tx.send(&data, vec![], vec![]).unwrap();
                    rx.recv().unwrap();
                }
                0
            });
        }
    }

    #[bench]
    fn transfer_data_00_1(b: &mut test::Bencher) {
        bench_transfer_data(b, 1);
    }
    #[bench]
    fn transfer_data_01_2(b: &mut test::Bencher) {
        bench_transfer_data(b, 2);
    }
    #[bench]
    fn transfer_data_02_4(b: &mut test::Bencher) {
        bench_transfer_data(b, 4);
    }
    #[bench]
    fn transfer_data_03_8(b: &mut test::Bencher) {
        bench_transfer_data(b, 8);
    }
    #[bench]
    fn transfer_data_04_16(b: &mut test::Bencher) {
        bench_transfer_data(b, 16);
    }
    #[bench]
    fn transfer_data_05_32(b: &mut test::Bencher) {
        bench_transfer_data(b, 32);
    }
    #[bench]
    fn transfer_data_06_64(b: &mut test::Bencher) {
        bench_transfer_data(b, 64);
    }
    #[bench]
    fn transfer_data_07_128(b: &mut test::Bencher) {
        bench_transfer_data(b, 128);
    }
    #[bench]
    fn transfer_data_08_256(b: &mut test::Bencher) {
        bench_transfer_data(b, 256);
    }
    #[bench]
    fn transfer_data_09_512(b: &mut test::Bencher) {
        bench_transfer_data(b, 512);
    }
    #[bench]
    fn transfer_data_10_1k(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024);
    }
    #[bench]
    fn transfer_data_11_2k(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024);
    }
    #[bench]
    fn transfer_data_12_4k(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024);
    }
    #[bench]
    fn transfer_data_13_8k(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024);
    }
    #[bench]
    fn transfer_data_14_16k(b: &mut test::Bencher) {
        bench_transfer_data(b, 16 * 1024);
    }
    #[bench]
    fn transfer_data_15_32k(b: &mut test::Bencher) {
        bench_transfer_data(b, 32 * 1024);
    }
    #[bench]
    fn transfer_data_16_64k(b: &mut test::Bencher) {
        bench_transfer_data(b, 64 * 1024);
    }
    #[bench]
    fn transfer_data_17_128k(b: &mut test::Bencher) {
        bench_transfer_data(b, 128 * 1024);
    }
    #[bench]
    fn transfer_data_18_256k(b: &mut test::Bencher) {
        bench_transfer_data(b, 256 * 1024);
    }
    #[bench]
    fn transfer_data_19_512k(b: &mut test::Bencher) {
        bench_transfer_data(b, 512 * 1024);
    }
    #[bench]
    fn transfer_data_20_1m(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_21_2m(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_22_4m(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_23_8m(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024 * 1024);
    }
}

mod ipc {
    mod receiver_set {
        use ipc_channel::ipc::{self, IpcReceiverSet};
        use test;

        // Benchmark selecting over a set of `n` receivers,
        // with `to_send` of them actually having pending data.
        fn gen_select_test(b: &mut test::Bencher, to_send: usize, n: usize) -> () {
            let mut active = Vec::with_capacity(to_send);
            let mut dormant = Vec::with_capacity(n - to_send);
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..to_send {
                let (tx, rx) = ipc::channel().unwrap();
                rx_set.add(rx).unwrap();
                active.push(tx);
            }
            for _ in to_send..n {
                let (tx, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
                dormant.push(tx);
            }
            b.iter(|| {
                for tx in active.iter() {
                    tx.send(()).unwrap();
                }
                let mut received = 0;
                while received < to_send {
                    for result in rx_set.select().unwrap().into_iter() {
                        let (_, _) = result.unwrap();
                        received += 1;
                    }
                }
            });
        }

        #[bench]
        fn send_on_1_of_1(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 1);
        }

        #[bench]
        fn send_on_1_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 5);
        }

        #[bench]
        fn send_on_2_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 2, 5);
        }

        #[bench]
        fn send_on_5_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 5);
        }

        #[bench]
        fn send_on_1_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 20);
        }

        #[bench]
        fn send_on_5_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 20);
        }

        #[bench]
        fn send_on_20_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 20, 20);
        }

        #[bench]
        fn send_on_1_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 100);
        }

        #[bench]
        fn send_on_5_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 100);
        }
        #[bench]
        fn send_on_20_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 20, 100);
        }

        #[bench]
        fn send_on_100_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 100, 100);
        }

        fn create_set_of_n(n: usize) -> IpcReceiverSet {
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..n {
                let (_, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
            }
            rx_set
        }

        #[bench]
        fn create_and_destroy_empty_set(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(0);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_1(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(1);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_10(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(10);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_100(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(100);
            });
        }

        // Benchmark performance of removing closed receivers from set.
        // This also includes the time for adding receivers,
        // as there is no way to measure the removing in isolation.
        fn bench_remove_closed(b: &mut test::Bencher, n: usize) {
            b.iter(|| {
                let mut rx_set = create_set_of_n(n);

                let mut dropped_count = 0;
                while dropped_count < n {
                    // On `select()`, receivers with a "ClosedChannel" event will be closed,
                    // and automatically dropped from the set.
                    dropped_count += rx_set.select().unwrap().len();
                }
            });
        }

        #[bench]
        fn add_and_remove_1_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 1);
        }

        #[bench]
        fn add_and_remove_10_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 10);
        }

        #[bench]
        fn add_and_remove_100_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 100);
        }
    }
}
<commit_msg>Benchmarks: cleanup: Use more consistent naming for parametric benchmark<commit_after>#![feature(test)]

extern crate ipc_channel;
extern crate test;

/// Allows doing multiple inner iterations per bench.iter() run.
///
/// This is mostly to amortise the overhead of spawning a thread in the benchmark
/// when sending larger messages (that might be fragmented).
///
/// Note that you need to compensate the displayed results
/// for the proportionally longer runs yourself,
/// as the benchmark framework doesn't know about the inner iterations...
const ITERATIONS: usize = 1;

mod platform {
    extern crate crossbeam;

    use ipc_channel::platform;
    use ITERATIONS;

    use std::sync::{mpsc, Mutex};
    use test;

    #[bench]
    fn create_channel(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..ITERATIONS {
                platform::channel().unwrap();
            }
        });
    }

    fn bench_transfer_data(b: &mut test::Bencher, size: usize) {
        let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
        let (tx, rx) = platform::channel().unwrap();

        let (wait_tx, wait_rx) = mpsc::channel();
        let wait_rx = Mutex::new(wait_rx);

        if size > platform::OsIpcSender::get_max_fragment_size() {
            b.iter(|| {
                crossbeam::scope(|scope| {
                    let tx = tx.clone();
                    scope.spawn(|| {
                        let wait_rx = wait_rx.lock().unwrap();
                        let tx = tx;
                        for _ in 0..ITERATIONS {
                            tx.send(&data, vec![], vec![]).unwrap();
                            if ITERATIONS > 1 {
                                // Prevent beginning of the next send
                                // from overlapping with receive of last fragment,
                                // as otherwise results of runs with a large tail fragment
                                // are significantly skewed.
                                wait_rx.recv().unwrap();
                            }
                        }
                    });
                    for _ in 0..ITERATIONS {
                        rx.recv().unwrap();
                        if ITERATIONS > 1 {
                            wait_tx.send(()).unwrap();
                        }
                    }
                    // For reasons mysterious to me,
                    // not returning a value *from every branch*
                    // adds some 100 ns or so of overhead to all results --
                    // which is quite significant for very short tests...
                    0
                })
            });
        } else {
            b.iter(|| {
                for _ in 0..ITERATIONS {
                    tx.send(&data, vec![], vec![]).unwrap();
                    rx.recv().unwrap();
                }
                0
            });
        }
    }

    #[bench]
    fn transfer_data_00_1(b: &mut test::Bencher) {
        bench_transfer_data(b, 1);
    }
    #[bench]
    fn transfer_data_01_2(b: &mut test::Bencher) {
        bench_transfer_data(b, 2);
    }
    #[bench]
    fn transfer_data_02_4(b: &mut test::Bencher) {
        bench_transfer_data(b, 4);
    }
    #[bench]
    fn transfer_data_03_8(b: &mut test::Bencher) {
        bench_transfer_data(b, 8);
    }
    #[bench]
    fn transfer_data_04_16(b: &mut test::Bencher) {
        bench_transfer_data(b, 16);
    }
    #[bench]
    fn transfer_data_05_32(b: &mut test::Bencher) {
        bench_transfer_data(b, 32);
    }
    #[bench]
    fn transfer_data_06_64(b: &mut test::Bencher) {
        bench_transfer_data(b, 64);
    }
    #[bench]
    fn transfer_data_07_128(b: &mut test::Bencher) {
        bench_transfer_data(b, 128);
    }
    #[bench]
    fn transfer_data_08_256(b: &mut test::Bencher) {
        bench_transfer_data(b, 256);
    }
    #[bench]
    fn transfer_data_09_512(b: &mut test::Bencher) {
        bench_transfer_data(b, 512);
    }
    #[bench]
    fn transfer_data_10_1k(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024);
    }
    #[bench]
    fn transfer_data_11_2k(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024);
    }
    #[bench]
    fn transfer_data_12_4k(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024);
    }
    #[bench]
    fn transfer_data_13_8k(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024);
    }
    #[bench]
    fn transfer_data_14_16k(b: &mut test::Bencher) {
        bench_transfer_data(b, 16 * 1024);
    }
    #[bench]
    fn transfer_data_15_32k(b: &mut test::Bencher) {
        bench_transfer_data(b, 32 * 1024);
    }
    #[bench]
    fn transfer_data_16_64k(b: &mut test::Bencher) {
        bench_transfer_data(b, 64 * 1024);
    }
    #[bench]
    fn transfer_data_17_128k(b: &mut test::Bencher) {
        bench_transfer_data(b, 128 * 1024);
    }
    #[bench]
    fn transfer_data_18_256k(b: &mut test::Bencher) {
        bench_transfer_data(b, 256 * 1024);
    }
    #[bench]
    fn transfer_data_19_512k(b: &mut test::Bencher) {
        bench_transfer_data(b, 512 * 1024);
    }
    #[bench]
    fn transfer_data_20_1m(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_21_2m(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_22_4m(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_23_8m(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024 * 1024);
    }
}

mod ipc {
    mod receiver_set {
        use ipc_channel::ipc::{self, IpcReceiverSet};
        use test;

        // Benchmark selecting over a set of `n` receivers,
        // with `to_send` of them actually having pending data.
        fn bench_send_on_m_of_n(b: &mut test::Bencher, to_send: usize, n: usize) -> () {
            let mut active = Vec::with_capacity(to_send);
            let mut dormant = Vec::with_capacity(n - to_send);
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..to_send {
                let (tx, rx) = ipc::channel().unwrap();
                rx_set.add(rx).unwrap();
                active.push(tx);
            }
            for _ in to_send..n {
                let (tx, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
                dormant.push(tx);
            }
            b.iter(|| {
                for tx in active.iter() {
                    tx.send(()).unwrap();
                }
                let mut received = 0;
                while received < to_send {
                    for result in rx_set.select().unwrap().into_iter() {
                        let (_, _) = result.unwrap();
                        received += 1;
                    }
                }
            });
        }

        #[bench]
        fn send_on_1_of_1(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 1, 1);
        }

        #[bench]
        fn send_on_1_of_5(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 1, 5);
        }

        #[bench]
        fn send_on_2_of_5(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 2, 5);
        }

        #[bench]
        fn send_on_5_of_5(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 5, 5);
        }

        #[bench]
        fn send_on_1_of_20(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 1, 20);
        }

        #[bench]
        fn send_on_5_of_20(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 5, 20);
        }

        #[bench]
        fn send_on_20_of_20(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 20, 20);
        }

        #[bench]
        fn send_on_1_of_100(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 1, 100);
        }

        #[bench]
        fn send_on_5_of_100(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 5, 100);
        }
        #[bench]
        fn send_on_20_of_100(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 20, 100);
        }

        #[bench]
        fn send_on_100_of_100(b: &mut test::Bencher) -> () {
            bench_send_on_m_of_n(b, 100, 100);
        }

        fn create_set_of_n(n: usize) -> IpcReceiverSet {
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..n {
                let (_, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
            }
            rx_set
        }

        #[bench]
        fn create_and_destroy_empty_set(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(0);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_1(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(1);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_10(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(10);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_100(b: &mut test::Bencher) -> () {
            b.iter(|| {
                create_set_of_n(100);
            });
        }

        // Benchmark performance of removing closed receivers from set.
        // This also includes the time for adding receivers,
        // as there is no way to measure the removing in isolation.
        fn bench_remove_closed(b: &mut test::Bencher, n: usize) {
            b.iter(|| {
                let mut rx_set = create_set_of_n(n);

                let mut dropped_count = 0;
                while dropped_count < n {
                    // On `select()`, receivers with a "ClosedChannel" event will be closed,
                    // and automatically dropped from the set.
                    dropped_count += rx_set.select().unwrap().len();
                }
            });
        }

        #[bench]
        fn add_and_remove_1_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 1);
        }

        #[bench]
        fn add_and_remove_10_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 10);
        }

        #[bench]
        fn add_and_remove_100_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 100);
        }
    }
}
<|endoftext|>


<commit_before>#![feature(raw, test)]

extern crate pico_sys as pico;
extern crate httparse;

extern crate test;

const REQ: &'static [u8] = b"\
GET /wp-content/uploads/2010/03/hello-kitty-darth-vader-pink.jpg HTTP/1.1\r\n\
Host: www.kittyhell.com\r\n\
User-Agent: Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10.6; ja-JP-mac; rv:1.9.2.3) Gecko/20100401 Firefox/3.6.3 Pathtraq/0.9\r\n\
Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8\r\n\
Accept-Language: ja,en-us;q=0.7,en;q=0.3\r\n\
Accept-Encoding: gzip,deflate\r\n\
Accept-Charset: Shift_JIS,utf-8;q=0.7,*;q=0.7\r\n\
Keep-Alive: 115\r\n\
Connection: keep-alive\r\n\
Cookie: wp_ozh_wsa_visits=2; wp_ozh_wsa_visit_lasttime=xxxxxxxxxx; __utma=xxxxxxxxx.xxxxxxxxxx.xxxxxxxxxx.xxxxxxxxxx.xxxxxxxxxx.x; __utmz=xxxxxxxxx.xxxxxxxxxx.x.x.utmccn=(referral)|utmcsr=reader.livedoor.com|utmcct=/reader/|utmcmd=referral\r\n\r\n";

#[bench]
fn bench_pico(b: &mut test::Bencher) {
    use std::{mem, raw};

    unsafe fn slice_to_mut_pair<'a, T>(slice: &'a mut &[T]) -> (&'a mut *const T, &'a mut usize) {
        let raw = mem::transmute::<_, &mut raw::Slice<T>>(slice);
        (&mut raw.data, &mut raw.len)
    }

    #[repr(C)]
    #[derive(Clone, Copy)]
    struct Header<'a>(&'a [u8], &'a [u8]);

    #[repr(C)]
    struct Headers<'a>(&'a mut [Header<'a>]);
    let method = [0i8; 16];
    let path = [0i8; 16];
    let mut minor_version = 0;
    let mut h = [Header(&[], &[]); 16];
    let headers = Headers(&mut h);
    let prev_buf_len = 0;

    b.iter(|| {
        let ret = unsafe {
            pico::ffi::phr_parse_request(
                REQ.as_ptr() as *const _,
                REQ.len() as u64,
                &mut method.as_ptr(),
                &mut 16,
                &mut path.as_ptr(),
                &mut 16,
                &mut minor_version,
                mem::transmute::<*mut Header, *mut pico::ffi::phr_header>(headers.0.as_mut_ptr()),
                slice_to_mut_pair(&mut &*headers.0).1 as *mut usize as *mut _,
                prev_buf_len
            )
        };
        assert_eq!(ret, REQ.len() as i32);
    })
}

#[bench]
fn bench_httparse(b: &mut test::Bencher) {
    let mut headers = [httparse::Header{ name: "", value: &[] }; 16];
    let mut req = httparse::Request::new(&mut headers);
    b.iter(|| {
        assert_eq!(req.parse(REQ).unwrap(), httparse::Status::Complete(REQ.len()));
    })
}
<commit_msg>fix benchmarks for newest nightlies<commit_after>#![feature(test)]

extern crate pico_sys as pico;
extern crate httparse;

extern crate test;

const REQ: &'static [u8] = b"\
GET /wp-content/uploads/2010/03/hello-kitty-darth-vader-pink.jpg HTTP/1.1\r\n\
Host: www.kittyhell.com\r\n\
User-Agent: Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10.6; ja-JP-mac; rv:1.9.2.3) Gecko/20100401 Firefox/3.6.3 Pathtraq/0.9\r\n\
Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8\r\n\
Accept-Language: ja,en-us;q=0.7,en;q=0.3\r\n\
Accept-Encoding: gzip,deflate\r\n\
Accept-Charset: Shift_JIS,utf-8;q=0.7,*;q=0.7\r\n\
Keep-Alive: 115\r\n\
Connection: keep-alive\r\n\
Cookie: wp_ozh_wsa_visits=2; wp_ozh_wsa_visit_lasttime=xxxxxxxxxx; __utma=xxxxxxxxx.xxxxxxxxxx.xxxxxxxxxx.xxxxxxxxxx.xxxxxxxxxx.x; __utmz=xxxxxxxxx.xxxxxxxxxx.x.x.utmccn=(referral)|utmcsr=reader.livedoor.com|utmcct=/reader/|utmcmd=referral\r\n\r\n";

#[bench]
fn bench_pico(b: &mut test::Bencher) {
    use std::mem;

    #[repr(C)]
    #[derive(Clone, Copy)]
    struct Header<'a>(&'a [u8], &'a [u8]);

    #[repr(C)]
    struct Headers<'a>(&'a mut [Header<'a>]);
    let method = [0i8; 16];
    let path = [0i8; 16];
    let mut minor_version = 0;
    let mut h = [Header(&[], &[]); 16];
    let mut h_len = h.len();
    let headers = Headers(&mut h);
    let prev_buf_len = 0;

    b.iter(|| {
        let ret = unsafe {
            pico::ffi::phr_parse_request(
                REQ.as_ptr() as *const _,
                REQ.len() as u64,
                &mut method.as_ptr(),
                &mut 16,
                &mut path.as_ptr(),
                &mut 16,
                &mut minor_version,
                mem::transmute::<*mut Header, *mut pico::ffi::phr_header>(headers.0.as_mut_ptr()),
                &mut h_len as *mut usize as *mut _,
                prev_buf_len
            )
        };
        assert_eq!(ret, REQ.len() as i32);
    })
}

#[bench]
fn bench_httparse(b: &mut test::Bencher) {
    let mut headers = [httparse::Header{ name: "", value: &[] }; 16];
    let mut req = httparse::Request::new(&mut headers);
    b.iter(|| {
        assert_eq!(req.parse(REQ).unwrap(), httparse::Status::Complete(REQ.len()));
    })
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate criterion;
extern crate wkt;

use criterion::Criterion;

fn criterion_benchmark(c: &mut criterion::Criterion) {
    c.bench_function("parse small", |bencher| {
        let s = include_str!("./small.wkt");
        bencher.iter(|| {
            let _ = wkt::Wkt::from_str(s).unwrap();
        });
    });

    c.bench_function("parse big", |bencher| {
        let s = include_str!("./big.wkt");
        bencher.iter(|| {
            let _ = wkt::Wkt::from_str(s).unwrap();
        });
    });
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
<commit_msg>Fix benchmark<commit_after>#[macro_use]
extern crate criterion;
extern crate wkt;

use criterion::Criterion;

fn criterion_benchmark(c: &mut criterion::Criterion) {
    c.bench_function("parse small", |bencher| {
        let s = include_str!("./small.wkt");
        bencher.iter(|| {
            let _ = wkt::Wkt::<f64>::from_str(s).unwrap();
        });
    });

    c.bench_function("parse big", |bencher| {
        let s = include_str!("./big.wkt");
        bencher.iter(|| {
            let _ = wkt::Wkt::<f64>::from_str(s).unwrap();
        });
    });
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
<|endoftext|>


<commit_before>//! Support datetimes and timedeltas
//!
//! [The corresponding section][datetime] of the NumPy documentation contains more information.
//!
//! # Example
//!
//! ```
//! use numpy::{datetime::{units, Datetime, Timedelta}, PyArray1};
//! use pyo3::Python;
//! # use pyo3::types::PyDict;
//!
//! Python::with_gil(|py| {
//! #    let locals = py
//! #        .eval("{ 'np': __import__('numpy') }", None, None)
//! #        .unwrap()
//! #        .downcast::<PyDict>()
//! #        .unwrap();
//! #
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Datetime<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Datetime::<units::Days>::from(17_277)
//!     );
//!
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2022-03-29')]) - np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Timedelta<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Timedelta::<units::Days>::from(1_803)
//!     );
//! });
//! ```
//!
//! [datetime]: https://numpy.org/doc/stable/reference/arrays.datetime.html

use std::fmt;
use std::hash::Hash;
use std::marker::PhantomData;

use pyo3::Python;

use crate::dtype::{Element, PyArrayDescr};
use crate::npyffi::{PyArray_DatetimeDTypeMetaData, NPY_DATETIMEUNIT, NPY_TYPES};

/// Represents the [datetime units][datetime-units] supported by NumPy
///
/// [datetime-units]: https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units
pub trait Unit: Send + Sync + Clone + Copy + PartialEq + Eq + Hash + PartialOrd + Ord {
    /// The matching NumPy [datetime unit code][NPY_DATETIMEUNIT]
    ///
    /// [NPY_DATETIMEUNIT]: https://github.com/numpy/numpy/blob/4c60b3263ac50e5e72f6a909e156314fc3c9cba0/numpy/core/include/numpy/ndarraytypes.h#L276
    const UNIT: NPY_DATETIMEUNIT;

    /// The abbrevation used for debug formatting
    const ABBREV: &'static str;
}

macro_rules! define_units {
    ($($(#[$meta:meta])* $struct:ident => $unit:ident $abbrev:literal,)+) => {
        $(

        $(#[$meta])*
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
        pub struct $struct;

        impl Unit for $struct {
            const UNIT: NPY_DATETIMEUNIT = NPY_DATETIMEUNIT::$unit;

            const ABBREV: &'static str = $abbrev;
        }

        )+
    };
}

/// Predefined implementors of the [`Unit`] trait
pub mod units {
    use super::*;

    define_units!(
        #[doc = "Years, i.e. 12 months"]
        Years => NPY_FR_Y "a",
        #[doc = "Months, i.e. 30 days"]
        Months => NPY_FR_M "mo",
        #[doc = "Weeks, i.e. 7 days"]
        Weeks => NPY_FR_W "w",
        #[doc = "Days, i.e. 24 hours"]
        Days => NPY_FR_D "d",
        #[doc = "Hours, i.e. 60 minutes"]
        Hours => NPY_FR_h "h",
        #[doc = "Minutes, i.e. 60 seconds"]
        Minutes => NPY_FR_m "min",
        #[doc = "Seconds"]
        Seconds => NPY_FR_s "s",
        #[doc = "Milliseconds, i.e. 10^-3 seconds"]
        Milliseconds => NPY_FR_ms "ms",
        #[doc = "Microseconds, i.e. 10^-6 seconds"]
        Microseconds => NPY_FR_us "µs",
        #[doc = "Nanoseconds, i.e. 10^-9 seconds"]
        Nanoseconds => NPY_FR_ns "ns",
        #[doc = "Picoseconds, i.e. 10^-12 seconds"]
        Picoseconds => NPY_FR_ps "ps",
        #[doc = "Femtoseconds, i.e. 10^-15 seconds"]
        Femtoseconds => NPY_FR_fs "fs",
        #[doc = "Attoseconds, i.e. 10^-18 seconds"]
        Attoseconds => NPY_FR_as "as",
    );
}

/// Corresponds to the [`datetime64`][scalars-datetime64] scalar type
///
/// [scalars-datetime64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.datetime64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Datetime<U: Unit>(i64, PhantomData<U>);

impl<U: Unit> From<i64> for Datetime<U> {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl<U: Unit> From<Datetime<U>> for i64 {
    fn from(val: Datetime<U>) -> Self {
        val.0
    }
}

unsafe impl<U: Unit> Element for Datetime<U> {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        // FIXME(adamreichold): Memoize these via the Unit trait
        let dtype = PyArrayDescr::new_from_npy_type(py, NPY_TYPES::NPY_DATETIME);

        unsafe {
            set_unit(dtype, U::UNIT);
        }

        dtype
    }
}

impl<U: Unit> fmt::Debug for Datetime<U> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Datetime({} {})", self.0, U::ABBREV)
    }
}

/// Corresponds to the [`timedelta64`][scalars-datetime64] scalar type
///
/// [scalars-timedelta64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.timedelta64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Timedelta<U: Unit>(i64, PhantomData<U>);

impl<U: Unit> From<i64> for Timedelta<U> {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl<U: Unit> From<Timedelta<U>> for i64 {
    fn from(val: Timedelta<U>) -> Self {
        val.0
    }
}

unsafe impl<U: Unit> Element for Timedelta<U> {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        // FIXME(adamreichold): Memoize these via the Unit trait
        let dtype = PyArrayDescr::new_from_npy_type(py, NPY_TYPES::NPY_TIMEDELTA);

        unsafe {
            set_unit(dtype, U::UNIT);
        }

        dtype
    }
}

impl<U: Unit> fmt::Debug for Timedelta<U> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Timedelta({} {})", self.0, U::ABBREV)
    }
}

unsafe fn set_unit(dtype: &PyArrayDescr, unit: NPY_DATETIMEUNIT) {
    let metadata = &mut *((*dtype.as_dtype_ptr()).c_metadata as *mut PyArray_DatetimeDTypeMetaData);

    metadata.meta.base = unit;
    metadata.meta.num = 1;
}

#[cfg(test)]
mod tests {
    use super::*;

    use pyo3::{
        py_run,
        types::{PyDict, PyModule},
    };

    use crate::array::PyArray1;

    #[test]
    fn from_python_to_rust() {
        Python::with_gil(|py| {
            let locals = py
                .eval("{ 'np': __import__('numpy') }", None, None)
                .unwrap()
                .downcast::<PyDict>()
                .unwrap();

            let array = py
                .eval(
                    "np.array([np.datetime64('1970-01-01')])",
                    None,
                    Some(locals),
                )
                .unwrap()
                .downcast::<PyArray1<Datetime<units::Days>>>()
                .unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, 0);
        });
    }

    #[test]
    fn from_rust_to_python() {
        Python::with_gil(|py| {
            let array = PyArray1::<Timedelta<units::Minutes>>::zeros(py, 1, false);

            *array.readwrite().get_mut(0).unwrap() = Timedelta::<units::Minutes>::from(5);

            let np = py
                .eval("__import__('numpy')", None, None)
                .unwrap()
                .downcast::<PyModule>()
                .unwrap();

            py_run!(py, array np, "assert array.dtype == np.dtype('timedelta64[m]')");
            py_run!(py, array np, "assert array[0] == np.timedelta64(5, 'm')");
        });
    }

    #[test]
    fn debug_formatting() {
        assert_eq!(
            format!("{:?}", Datetime::<units::Days>::from(28)),
            "Datetime(28 d)"
        );

        assert_eq!(
            format!("{:?}", Timedelta::<units::Milliseconds>::from(160)),
            "Timedelta(160 ms)"
        );
    }

    #[test]
    fn unit_conversion() {
        #[track_caller]
        fn convert<S: Unit, D: Unit>(py: Python<'_>, expected_value: i64) {
            let array = PyArray1::<Timedelta<S>>::from_slice(py, &[Timedelta::<S>::from(1)]);
            let array = array.cast::<Timedelta<D>>(false).unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, expected_value);
        }

        Python::with_gil(|py| {
            convert::<units::Years, units::Days>(py, (97 + 400 * 365) / 400);
            convert::<units::Months, units::Days>(py, (97 + 400 * 365) / 400 / 12);

            convert::<units::Weeks, units::Seconds>(py, 7 * 24 * 60 * 60);
            convert::<units::Days, units::Seconds>(py, 24 * 60 * 60);
            convert::<units::Hours, units::Seconds>(py, 60 * 60);
            convert::<units::Minutes, units::Seconds>(py, 60);

            convert::<units::Seconds, units::Milliseconds>(py, 1_000);
            convert::<units::Seconds, units::Microseconds>(py, 1_000_000);
            convert::<units::Seconds, units::Nanoseconds>(py, 1_000_000_000);
            convert::<units::Seconds, units::Picoseconds>(py, 1_000_000_000_000);
            convert::<units::Seconds, units::Femtoseconds>(py, 1_000_000_000_000_000);

            convert::<units::Femtoseconds, units::Attoseconds>(py, 1_000);
        });
    }
}
<commit_msg>Memoize type descriptors for datetime/timedelta values on a per-unit basis.<commit_after>//! Support datetimes and timedeltas
//!
//! [The corresponding section][datetime] of the NumPy documentation contains more information.
//!
//! # Example
//!
//! ```
//! use numpy::{datetime::{units, Datetime, Timedelta}, PyArray1};
//! use pyo3::Python;
//! # use pyo3::types::PyDict;
//!
//! Python::with_gil(|py| {
//! #    let locals = py
//! #        .eval("{ 'np': __import__('numpy') }", None, None)
//! #        .unwrap()
//! #        .downcast::<PyDict>()
//! #        .unwrap();
//! #
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Datetime<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Datetime::<units::Days>::from(17_277)
//!     );
//!
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2022-03-29')]) - np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Timedelta<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Timedelta::<units::Days>::from(1_803)
//!     );
//! });
//! ```
//!
//! [datetime]: https://numpy.org/doc/stable/reference/arrays.datetime.html

use std::cell::UnsafeCell;
use std::collections::hash_map::Entry;
use std::fmt;
use std::hash::Hash;
use std::marker::PhantomData;

use ahash::AHashMap;
use pyo3::{Py, Python};

use crate::dtype::{Element, PyArrayDescr};
use crate::npyffi::{PyArray_DatetimeDTypeMetaData, NPY_DATETIMEUNIT, NPY_TYPES};

/// Represents the [datetime units][datetime-units] supported by NumPy
///
/// [datetime-units]: https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units
pub trait Unit: Send + Sync + Clone + Copy + PartialEq + Eq + Hash + PartialOrd + Ord {
    /// The matching NumPy [datetime unit code][NPY_DATETIMEUNIT]
    ///
    /// [NPY_DATETIMEUNIT]: https://github.com/numpy/numpy/blob/4c60b3263ac50e5e72f6a909e156314fc3c9cba0/numpy/core/include/numpy/ndarraytypes.h#L276
    const UNIT: NPY_DATETIMEUNIT;

    /// The abbrevation used for debug formatting
    const ABBREV: &'static str;
}

macro_rules! define_units {
    ($($(#[$meta:meta])* $struct:ident => $unit:ident $abbrev:literal,)+) => {
        $(

        $(#[$meta])*
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
        pub struct $struct;

        impl Unit for $struct {
            const UNIT: NPY_DATETIMEUNIT = NPY_DATETIMEUNIT::$unit;

            const ABBREV: &'static str = $abbrev;
        }

        )+
    };
}

/// Predefined implementors of the [`Unit`] trait
pub mod units {
    use super::*;

    define_units!(
        #[doc = "Years, i.e. 12 months"]
        Years => NPY_FR_Y "a",
        #[doc = "Months, i.e. 30 days"]
        Months => NPY_FR_M "mo",
        #[doc = "Weeks, i.e. 7 days"]
        Weeks => NPY_FR_W "w",
        #[doc = "Days, i.e. 24 hours"]
        Days => NPY_FR_D "d",
        #[doc = "Hours, i.e. 60 minutes"]
        Hours => NPY_FR_h "h",
        #[doc = "Minutes, i.e. 60 seconds"]
        Minutes => NPY_FR_m "min",
        #[doc = "Seconds"]
        Seconds => NPY_FR_s "s",
        #[doc = "Milliseconds, i.e. 10^-3 seconds"]
        Milliseconds => NPY_FR_ms "ms",
        #[doc = "Microseconds, i.e. 10^-6 seconds"]
        Microseconds => NPY_FR_us "µs",
        #[doc = "Nanoseconds, i.e. 10^-9 seconds"]
        Nanoseconds => NPY_FR_ns "ns",
        #[doc = "Picoseconds, i.e. 10^-12 seconds"]
        Picoseconds => NPY_FR_ps "ps",
        #[doc = "Femtoseconds, i.e. 10^-15 seconds"]
        Femtoseconds => NPY_FR_fs "fs",
        #[doc = "Attoseconds, i.e. 10^-18 seconds"]
        Attoseconds => NPY_FR_as "as",
    );
}

/// Corresponds to the [`datetime64`][scalars-datetime64] scalar type
///
/// [scalars-datetime64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.datetime64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Datetime<U: Unit>(i64, PhantomData<U>);

impl<U: Unit> From<i64> for Datetime<U> {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl<U: Unit> From<Datetime<U>> for i64 {
    fn from(val: Datetime<U>) -> Self {
        val.0
    }
}

unsafe impl<U: Unit> Element for Datetime<U> {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        static DTYPES: TypeDescriptors = unsafe { TypeDescriptors::new(NPY_TYPES::NPY_DATETIME) };

        DTYPES.from_unit(py, U::UNIT)
    }
}

impl<U: Unit> fmt::Debug for Datetime<U> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Datetime({} {})", self.0, U::ABBREV)
    }
}

/// Corresponds to the [`timedelta64`][scalars-datetime64] scalar type
///
/// [scalars-timedelta64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.timedelta64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Timedelta<U: Unit>(i64, PhantomData<U>);

impl<U: Unit> From<i64> for Timedelta<U> {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl<U: Unit> From<Timedelta<U>> for i64 {
    fn from(val: Timedelta<U>) -> Self {
        val.0
    }
}

unsafe impl<U: Unit> Element for Timedelta<U> {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        static DTYPES: TypeDescriptors = unsafe { TypeDescriptors::new(NPY_TYPES::NPY_TIMEDELTA) };

        DTYPES.from_unit(py, U::UNIT)
    }
}

impl<U: Unit> fmt::Debug for Timedelta<U> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Timedelta({} {})", self.0, U::ABBREV)
    }
}

struct TypeDescriptors {
    npy_type: NPY_TYPES,
    dtypes: UnsafeCell<Option<AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>>>>,
}

unsafe impl Sync for TypeDescriptors {}

impl TypeDescriptors {
    /// `npy_type` must be either `NPY_DATETIME` or `NPY_TIMEDELTA`.
    const unsafe fn new(npy_type: NPY_TYPES) -> Self {
        Self {
            npy_type,
            dtypes: UnsafeCell::new(None),
        }
    }

    #[allow(clippy::mut_from_ref)]
    unsafe fn get(&self) -> &mut AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>> {
        (*self.dtypes.get()).get_or_insert_with(AHashMap::new)
    }

    #[allow(clippy::wrong_self_convention)]
    fn from_unit<'py>(&'py self, py: Python<'py>, unit: NPY_DATETIMEUNIT) -> &'py PyArrayDescr {
        // SAFETY: We hold the GIL and we do not call into user code which might re-enter.
        let dtypes = unsafe { self.get() };

        match dtypes.entry(unit) {
            Entry::Occupied(entry) => entry.into_mut().as_ref(py),
            Entry::Vacant(entry) => {
                let dtype = PyArrayDescr::new_from_npy_type(py, self.npy_type);

                // SAFETY: `self.npy_type` is either `NPY_DATETIME` or `NPY_TIMEDELTA` which implies the type of `c_metadata`.
                unsafe {
                    let metadata = &mut *((*dtype.as_dtype_ptr()).c_metadata
                        as *mut PyArray_DatetimeDTypeMetaData);

                    metadata.meta.base = unit;
                    metadata.meta.num = 1;
                }

                entry.insert(dtype.into()).as_ref(py)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use pyo3::{
        py_run,
        types::{PyDict, PyModule},
    };

    use crate::array::PyArray1;

    #[test]
    fn from_python_to_rust() {
        Python::with_gil(|py| {
            let locals = py
                .eval("{ 'np': __import__('numpy') }", None, None)
                .unwrap()
                .downcast::<PyDict>()
                .unwrap();

            let array = py
                .eval(
                    "np.array([np.datetime64('1970-01-01')])",
                    None,
                    Some(locals),
                )
                .unwrap()
                .downcast::<PyArray1<Datetime<units::Days>>>()
                .unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, 0);
        });
    }

    #[test]
    fn from_rust_to_python() {
        Python::with_gil(|py| {
            let array = PyArray1::<Timedelta<units::Minutes>>::zeros(py, 1, false);

            *array.readwrite().get_mut(0).unwrap() = Timedelta::<units::Minutes>::from(5);

            let np = py
                .eval("__import__('numpy')", None, None)
                .unwrap()
                .downcast::<PyModule>()
                .unwrap();

            py_run!(py, array np, "assert array.dtype == np.dtype('timedelta64[m]')");
            py_run!(py, array np, "assert array[0] == np.timedelta64(5, 'm')");
        });
    }

    #[test]
    fn debug_formatting() {
        assert_eq!(
            format!("{:?}", Datetime::<units::Days>::from(28)),
            "Datetime(28 d)"
        );

        assert_eq!(
            format!("{:?}", Timedelta::<units::Milliseconds>::from(160)),
            "Timedelta(160 ms)"
        );
    }

    #[test]
    fn unit_conversion() {
        #[track_caller]
        fn convert<S: Unit, D: Unit>(py: Python<'_>, expected_value: i64) {
            let array = PyArray1::<Timedelta<S>>::from_slice(py, &[Timedelta::<S>::from(1)]);
            let array = array.cast::<Timedelta<D>>(false).unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, expected_value);
        }

        Python::with_gil(|py| {
            convert::<units::Years, units::Days>(py, (97 + 400 * 365) / 400);
            convert::<units::Months, units::Days>(py, (97 + 400 * 365) / 400 / 12);

            convert::<units::Weeks, units::Seconds>(py, 7 * 24 * 60 * 60);
            convert::<units::Days, units::Seconds>(py, 24 * 60 * 60);
            convert::<units::Hours, units::Seconds>(py, 60 * 60);
            convert::<units::Minutes, units::Seconds>(py, 60);

            convert::<units::Seconds, units::Milliseconds>(py, 1_000);
            convert::<units::Seconds, units::Microseconds>(py, 1_000_000);
            convert::<units::Seconds, units::Nanoseconds>(py, 1_000_000_000);
            convert::<units::Seconds, units::Picoseconds>(py, 1_000_000_000_000);
            convert::<units::Seconds, units::Femtoseconds>(py, 1_000_000_000_000_000);

            convert::<units::Femtoseconds, units::Attoseconds>(py, 1_000);
        });
    }
}
<|endoftext|>


<commit_before>use std::ops::{Mul, Add};

use rand::{Rand, Rng, Closed01};
use num::{Float, Zero};

use rblas::{Gemm, Matrix};
use rblas::attribute::Transpose;

use ndarray::{ArrayBase, DataOwned, DataMut};
use ndarray::blas::{BlasArrayViewMut, AsBlas};

use helpers::random01;

pub trait ShapeAsTuple<T> {
    fn shape_as_tuple(&self) -> T;
}

pub type FloatT = f32;

impl<T> ShapeAsTuple<(usize, usize)> for ArrayBase<Vec<T>, (usize, usize)> {
    #[inline]
    fn shape_as_tuple(&self) -> (usize, usize) {
        (self.shape()[0], self.shape()[1])
    }
}

pub type Ix = (usize, usize);
pub type MyMatrix = ArrayBase<Vec<FloatT>, Ix>;

pub struct OrthogonalNMFBlas {
    pub hidden: MyMatrix,
    pub weights: MyMatrix,

    // these hold temporary results during an iteration.
    // kept in this struct to prevent unnecessary memory allocations.
    pub weights_multiplier: MyMatrix,
    pub weights_divisor: MyMatrix,
    pub weights_divisor_reconstruction: MyMatrix,

    pub hidden_multiplier: MyMatrix,
    pub hidden_divisor: MyMatrix,
    pub hidden_divisor_partial: MyMatrix,

    pub gamma: MyMatrix,
}

impl OrthogonalNMFBlas
    // where FloatT: 'a + Float + Mul + Zero + Clone + Gemm,
          // ArrayBase<Vec<FloatT>, (usize, usize)>: 'a + AsBlas<FloatT, ArrayBase<Vec<FloatT>, (usize, usize)>, (usize, usize)> + DataOwned<Elem=FloatT> + DataMut<Elem=FloatT>,
          // BlasArrayViewMut<'a, FloatT, (usize, usize)>: Matrix<FloatT>,
          // Closed01<FloatT>: Rand
{
    pub fn new_random01<R: Rng>(nhidden: usize, nobserved: usize, nsamples: usize, rng: &mut R) -> OrthogonalNMFBlas {
        let mut hidden = MyMatrix::zeros((nhidden, nobserved));
        for x in hidden.iter_mut() {
            *x = random01(rng);
        }

        let mut weights = MyMatrix::zeros((nsamples, nhidden));
        for x in weights.iter_mut() {
            *x = random01(rng);
        }

        Self::new(hidden, weights)
    }

    pub fn new(hidden: MyMatrix, weights: MyMatrix) -> OrthogonalNMFBlas {
        let weights_shape = weights.shape_as_tuple();
        let hidden_shape = hidden.shape_as_tuple();

        let nhidden = hidden_shape.0;
        let nobserved = hidden_shape.1;
        let nsamples = weights_shape.0;

        let samples_shape = (nsamples, nobserved);

        OrthogonalNMFBlas {
            hidden: hidden,
            weights: weights,

            weights_multiplier: MyMatrix::zeros(weights_shape),
            weights_divisor: MyMatrix::zeros(weights_shape),
            weights_divisor_reconstruction: MyMatrix::zeros(samples_shape),

            hidden_multiplier: MyMatrix::zeros(hidden_shape),
            hidden_divisor: MyMatrix::zeros(hidden_shape),
            hidden_divisor_partial: MyMatrix::zeros((nhidden, nhidden)),

            gamma: OrthogonalNMFBlas::gamma(nhidden),
        }
    }

    /// returns the number of observed variables
    #[inline]
    pub fn nobserved(&self) -> usize {
        self.hidden.shape()[1]
    }

    /// returns the number of hidden variables
    #[inline]
    pub fn nhidden(&self) -> usize {
        self.hidden.shape()[0]
    }

    /// returns the number of data points
    #[inline]
    pub fn nsamples(&self) -> usize {
        self.weights.shape()[0]
    }

    pub fn samples_shape(&self) -> (usize, usize) {
        (self.nsamples(), self.nobserved())
    }

    /// gamma is a symetric matrix with diagonal elements equal to zero
    /// and other elements equal to 1
    pub fn gamma(size: usize) -> MyMatrix {
        let mut gamma = MyMatrix::from_elem((size, size), 1.);
        for x in gamma.diag_mut().iter_mut() {
            *x = 0.
        }
        gamma
    }

    /// `samples * hidden.transpose()`
    #[inline]
    pub fn update_weights_multiplier(&mut self, samples: &mut MyMatrix) {
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &samples.blas(),
            Transpose::Trans, &self.hidden.blas(),
            &0.,
            &mut self.weights_multiplier.blas());

    }

    /// `weights * hidden * hidden.transpose()`
    #[inline]
    pub fn update_weights_divisor(&mut self) {
        // TODO weights * hidden which we compute first is really large
        // we could first compute partial = hidden * hidden.transpose()
        // partial.shape() = (nhidden, nhidden)
        // and then in the second step compute
        // weights * partial.
        // however since we can't alias hidden we need to copy
        // it for this.
        // the question is whether the copying is faster than
        // the larger matrix multiplication. i guess it is.
        // TODO multiplying a matrix by its transpose is symetric
        // can we exploit that ?
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &self.weights.blas(),
            Transpose::NoTrans, &self.hidden.blas(),
            &0.,
            &mut self.weights_divisor_reconstruction.blas());
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &self.weights_divisor_reconstruction.blas(),
            Transpose::Trans, &self.hidden.blas(),
            &0.,
            &mut self.weights_divisor.blas());
    }

    /// `weights.transpose() * samples`
    #[inline]
    pub fn update_hidden_multiplier(&mut self, samples: &mut MyMatrix) {
        Gemm::gemm(
            &1.,
            Transpose::Trans, &self.weights.blas(),
            Transpose::NoTrans, &samples.blas(),
            &0.,
            &mut self.hidden_multiplier.blas());
    }

    /// `weights.transpose() * weights * hidden + alpha * gamma * hidden`
    #[inline]
    pub fn update_hidden_divisor(&mut self, alpha: FloatT) {
        // we must make a copy here because we need two mutable
        // references to weights at the same time for .blas()
        // weights_copy.clone_from(&weights);
        // TODO this will be solved once
        // https://github.com/bluss/rust-ndarray/pull/40
        // is merged.
        let mut weights_copy = self.weights.clone();
        Gemm::gemm(
            &1.,
            Transpose::Trans, &self.weights.blas(),
            Transpose::NoTrans, &weights_copy.blas(),
            &0.,
            &mut self.hidden_divisor_partial.blas());
        Gemm::gemm(
            &alpha,
            Transpose::NoTrans, &self.gamma.blas(),
            Transpose::NoTrans, &self.hidden.blas(),
            &0.,
            &mut self.hidden_divisor.blas());
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &self.hidden_divisor_partial.blas(),
            Transpose::NoTrans, &self.hidden.blas(),
            // add to previous contents of self.hidden_divisor
            &1.,
            &mut self.hidden_divisor.blas());
    }

    /// `result(i,j) <- result(i,j) * multiplier(i,j) / divisor(i,j)`
    #[inline]
    pub fn update_from_multiplier_and_divisor(
        multiplier: &MyMatrix,
        divisor: &MyMatrix,
        result: &mut MyMatrix,
    ) {
        let shape = result.shape_as_tuple();
        assert_eq!(shape, multiplier.shape_as_tuple());
        assert_eq!(shape, divisor.shape_as_tuple());

        for row in 0..shape.0 {
            for col in 0..shape.1 {
                let index = (row, col);
                let mut div = unsafe { divisor.uget(index).clone() };
                // if we have any 0 in any of the matrixes
                // self.weights or self.hidden then
                // divisor will be 0.
                // we can't divide by 0.
                // so we change them to the min positive value.
                if FloatT::zero() == div {
                    div = FloatT::min_positive_value();
                }
                unsafe {
                    *result.uget_mut(index) *= multiplier.uget(index) / div;
                }
            }
        }
    }

    // TODO consider calling this something like iteration_step
    // TODO how many iterations ?
    // TODO compare this to the seoung solution
    /// does one iteration step.
    /// `weights` and `hidden` get better and better with each iteration.
    /// usually around `10000` iterations are required.
    /// `samples` contains one observed per column, one sample per row.
    pub fn iterate(&mut self, alpha: FloatT, samples: &mut MyMatrix) {
        assert_eq!(samples.shape_as_tuple(), self.samples_shape());

        // weights_multiplier <- samples * hidden.transpose()
        self.update_weights_multiplier(samples);
        // weights_divisor <- weights * hidden * hidden.transpose()
        self.update_weights_divisor();
        // hidden_multiplier <- weights.transpose() * samples
        self.update_hidden_multiplier(samples);
        // hidden_divisor <-
        // weights.transpose() * weights * hidden + alpha * gamma * hidden
        self.update_hidden_divisor(alpha);

        // weights(i,j) <-
        // weights(i,j) * weights_multiplier(i,j) / weights_divisor(i,j)
        OrthogonalNMFBlas::update_from_multiplier_and_divisor(
            &self.weights_multiplier,
            &self.weights_divisor,
            &mut self.weights);
        // hidden(i,j) <-
        // hidden(i,j) * hidden_multiplier(i,j) / hidden_divisor(i,j)
        OrthogonalNMFBlas::update_from_multiplier_and_divisor(
            &self.hidden_multiplier,
            &self.hidden_divisor,
            &mut self.hidden);
    }
}
<commit_msg>improve comments<commit_after>use std::ops::{Mul, Add};

use rand::{Rand, Rng, Closed01};
use num::{Float, Zero};

use rblas::{Gemm, Matrix};
use rblas::attribute::Transpose;

use ndarray::{ArrayBase, DataOwned, DataMut};
use ndarray::blas::{BlasArrayViewMut, AsBlas};

use helpers::random01;

pub trait ShapeAsTuple<T> {
    fn shape_as_tuple(&self) -> T;
}

pub type FloatT = f32;

impl<T> ShapeAsTuple<(usize, usize)> for ArrayBase<Vec<T>, (usize, usize)> {
    #[inline]
    fn shape_as_tuple(&self) -> (usize, usize) {
        (self.shape()[0], self.shape()[1])
    }
}

pub type Ix = (usize, usize);
pub type MyMatrix = ArrayBase<Vec<FloatT>, Ix>;

pub struct OrthogonalNMFBlas {
    pub hidden: MyMatrix,
    pub weights: MyMatrix,

    // these hold temporary results during an iteration.
    // kept in this struct to prevent unnecessary memory allocations.
    pub weights_multiplier: MyMatrix,
    pub weights_divisor: MyMatrix,
    pub weights_divisor_reconstruction: MyMatrix,

    pub hidden_multiplier: MyMatrix,
    pub hidden_divisor: MyMatrix,
    pub hidden_divisor_partial: MyMatrix,

    pub gamma: MyMatrix,
}

impl OrthogonalNMFBlas
    // where FloatT: 'a + Float + Mul + Zero + Clone + Gemm,
          // ArrayBase<Vec<FloatT>, (usize, usize)>: 'a + AsBlas<FloatT, ArrayBase<Vec<FloatT>, (usize, usize)>, (usize, usize)> + DataOwned<Elem=FloatT> + DataMut<Elem=FloatT>,
          // BlasArrayViewMut<'a, FloatT, (usize, usize)>: Matrix<FloatT>,
          // Closed01<FloatT>: Rand
{
    pub fn new_random01<R: Rng>(nhidden: usize, nobserved: usize, nsamples: usize, rng: &mut R) -> OrthogonalNMFBlas {
        let mut hidden = MyMatrix::zeros((nhidden, nobserved));
        for x in hidden.iter_mut() {
            *x = random01(rng);
        }

        let mut weights = MyMatrix::zeros((nsamples, nhidden));
        for x in weights.iter_mut() {
            *x = random01(rng);
        }

        Self::new(hidden, weights)
    }

    pub fn new(hidden: MyMatrix, weights: MyMatrix) -> OrthogonalNMFBlas {
        let weights_shape = weights.shape_as_tuple();
        let hidden_shape = hidden.shape_as_tuple();

        let nhidden = hidden_shape.0;
        let nobserved = hidden_shape.1;
        let nsamples = weights_shape.0;

        let samples_shape = (nsamples, nobserved);

        OrthogonalNMFBlas {
            hidden: hidden,
            weights: weights,

            weights_multiplier: MyMatrix::zeros(weights_shape),
            weights_divisor: MyMatrix::zeros(weights_shape),
            weights_divisor_reconstruction: MyMatrix::zeros(samples_shape),

            hidden_multiplier: MyMatrix::zeros(hidden_shape),
            hidden_divisor: MyMatrix::zeros(hidden_shape),
            hidden_divisor_partial: MyMatrix::zeros((nhidden, nhidden)),

            gamma: OrthogonalNMFBlas::gamma(nhidden),
        }
    }

    /// returns the number of observed variables
    #[inline]
    pub fn nobserved(&self) -> usize {
        self.hidden.shape()[1]
    }

    /// returns the number of hidden variables
    #[inline]
    pub fn nhidden(&self) -> usize {
        self.hidden.shape()[0]
    }

    /// returns the number of data points
    #[inline]
    pub fn nsamples(&self) -> usize {
        self.weights.shape()[0]
    }

    pub fn samples_shape(&self) -> (usize, usize) {
        (self.nsamples(), self.nobserved())
    }

    /// gamma is a symetric matrix with diagonal elements equal to zero
    /// and other elements equal to 1
    pub fn gamma(size: usize) -> MyMatrix {
        let mut gamma = MyMatrix::from_elem((size, size), 1.);
        for x in gamma.diag_mut().iter_mut() {
            *x = 0.
        }
        gamma
    }

    /// `samples * hidden.transpose()`
    #[inline]
    pub fn update_weights_multiplier(&mut self, samples: &mut MyMatrix) {
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &samples.blas(),
            Transpose::Trans, &self.hidden.blas(),
            &0.,
            &mut self.weights_multiplier.blas());

    }

    // TODO this is the slowest
    /// `weights * hidden * hidden.transpose()`
    #[inline]
    pub fn update_weights_divisor(&mut self) {
        // TODO weights * hidden which we compute first is really large
        // we could first compute partial = hidden * hidden.transpose()
        // partial.shape() = (nhidden, nhidden)
        // and then in the second step compute
        // weights * partial.
        // however since we can't alias hidden we need to copy
        // it for this.
        // the question is whether the copying is faster than
        // the larger matrix multiplication. i guess it is.
        // TODO multiplying a matrix by its transpose is symetric
        // can we exploit that ?
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &self.weights.blas(),
            Transpose::NoTrans, &self.hidden.blas(),
            &0.,
            &mut self.weights_divisor_reconstruction.blas());
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &self.weights_divisor_reconstruction.blas(),
            Transpose::Trans, &self.hidden.blas(),
            &0.,
            &mut self.weights_divisor.blas());
    }

    /// `weights.transpose() * samples`
    #[inline]
    pub fn update_hidden_multiplier(&mut self, samples: &mut MyMatrix) {
        Gemm::gemm(
            &1.,
            Transpose::Trans, &self.weights.blas(),
            Transpose::NoTrans, &samples.blas(),
            &0.,
            &mut self.hidden_multiplier.blas());
    }

    // surprisingly this is the fastest
    /// `weights.transpose() * weights * hidden + alpha * gamma * hidden`
    #[inline]
    pub fn update_hidden_divisor(&mut self, alpha: FloatT) {
        // TODO benchmark whether using clone_from is faster
        let mut weights_copy = self.weights.clone();
        // partial <- weights.transpose() * weights
        Gemm::gemm(
            &1.,
            Transpose::Trans, &self.weights.blas(),
            Transpose::NoTrans, &weights_copy.blas(),
            &0.,
            &mut self.hidden_divisor_partial.blas());
        // partial.shape() = (nhidden, nhidden) very small
        // divisor <- alpha * gamma * hidden
        Gemm::gemm(
            &alpha,
            Transpose::NoTrans, &self.gamma.blas(),
            Transpose::NoTrans, &self.hidden.blas(),
            &0.,
            &mut self.hidden_divisor.blas());
        // divisor <- divisor + partial * hidden
        Gemm::gemm(
            &1.,
            Transpose::NoTrans, &self.hidden_divisor_partial.blas(),
            Transpose::NoTrans, &self.hidden.blas(),
            // add to previous contents of self.hidden_divisor
            &1.,
            &mut self.hidden_divisor.blas());
    }

    /// `result(i,j) <- result(i,j) * multiplier(i,j) / divisor(i,j)`
    #[inline]
    pub fn update_from_multiplier_and_divisor(
        multiplier: &MyMatrix,
        divisor: &MyMatrix,
        result: &mut MyMatrix,
    ) {
        let shape = result.shape_as_tuple();
        assert_eq!(shape, multiplier.shape_as_tuple());
        assert_eq!(shape, divisor.shape_as_tuple());

        for row in 0..shape.0 {
            for col in 0..shape.1 {
                let index = (row, col);
                let mut div = unsafe { divisor.uget(index).clone() };
                // if we have any 0 in any of the matrixes
                // self.weights or self.hidden then
                // divisor will be 0.
                // we can't divide by 0.
                // so we change them to the min positive value.
                if FloatT::zero() == div {
                    div = FloatT::min_positive_value();
                }
                unsafe {
                    *result.uget_mut(index) *= multiplier.uget(index) / div;
                }
            }
        }
    }

    // TODO consider calling this something like iteration_step
    // TODO how many iterations ?
    // TODO compare this to the seoung solution
    /// does one iteration step.
    /// `weights` and `hidden` get better and better with each iteration.
    /// usually around `10000` iterations are required.
    /// `samples` contains one observed per column, one sample per row.
    pub fn iterate(&mut self, alpha: FloatT, samples: &mut MyMatrix) {
        assert_eq!(samples.shape_as_tuple(), self.samples_shape());

        // weights_multiplier <- samples * hidden.transpose()
        self.update_weights_multiplier(samples);
        // weights_divisor <- weights * hidden * hidden.transpose()
        self.update_weights_divisor();
        // hidden_multiplier <- weights.transpose() * samples
        self.update_hidden_multiplier(samples);
        // hidden_divisor <-
        // weights.transpose() * weights * hidden + alpha * gamma * hidden
        self.update_hidden_divisor(alpha);

        // weights(i,j) <-
        // weights(i,j) * weights_multiplier(i,j) / weights_divisor(i,j)
        OrthogonalNMFBlas::update_from_multiplier_and_divisor(
            &self.weights_multiplier,
            &self.weights_divisor,
            &mut self.weights);
        // hidden(i,j) <-
        // hidden(i,j) * hidden_multiplier(i,j) / hidden_divisor(i,j)
        OrthogonalNMFBlas::update_from_multiplier_and_divisor(
            &self.hidden_multiplier,
            &self.hidden_divisor,
            &mut self.hidden);
    }
}
<|endoftext|>