<commit_before>use cairo::{self, ImageSurface};
use markup5ever::{expanded_name, local_name, namespace_url, ns, QualName};
use nalgebra::{Matrix3, Matrix4x5, Matrix5, Vector5};

use crate::drawing_ctx::DrawingCtx;
use crate::error::{AttributeResultExt, NodeError};
use crate::node::{NodeResult, NodeTrait, RsvgNode};
use crate::number_list::{NumberList, NumberListError, NumberListLength};
use crate::parsers;
use crate::property_bag::PropertyBag;
use crate::surface_utils::{
    iterators::Pixels, shared_surface::SharedImageSurface, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;

use super::context::{FilterContext, FilterOutput, FilterResult};
use super::{FilterEffect, FilterError, PrimitiveWithInput};

/// Color matrix operation types.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum OperationType {
    Matrix,
    Saturate,
    HueRotate,
    LuminanceToAlpha,
}

/// The `feColorMatrix` filter primitive.
pub struct FeColorMatrix {
    base: PrimitiveWithInput,
    matrix: Matrix5<f64>,
}

impl Default for FeColorMatrix {
    /// Constructs a new `ColorMatrix` with empty properties.
    #[inline]
    fn default() -> FeColorMatrix {
        FeColorMatrix {
            base: PrimitiveWithInput::new::<Self>(),
            matrix: Matrix5::identity(),
        }
    }
}

impl NodeTrait for FeColorMatrix {
    impl_node_as_filter_effect!();

    fn set_atts(&mut self, parent: Option<&RsvgNode>, pbag: &PropertyBag<'_>) -> NodeResult {
        self.base.set_atts(parent, pbag)?;

        // First, determine the operation type.
        let mut operation_type = OperationType::Matrix;
        for (attr, value) in pbag
            .iter()
            .filter(|(attr, _)| attr.expanded() == expanded_name!(svg "type"))
        {
            operation_type = OperationType::parse(attr, value)?;
        }

        // Now read the matrix correspondingly.
        // LuminanceToAlpha doesn't accept any matrix.
        if operation_type == OperationType::LuminanceToAlpha {
            self.matrix = {
                #[cfg_attr(rustfmt, rustfmt_skip)]
                Matrix5::new(
                    0.0,    0.0,    0.0,    0.0, 0.0,
                    0.0,    0.0,    0.0,    0.0, 0.0,
                    0.0,    0.0,    0.0,    0.0, 0.0,
                    0.2125, 0.7154, 0.0721, 0.0, 0.0,
                    0.0,    0.0,    0.0,    0.0, 1.0,
                )
            };
        } else {
            for (attr, value) in pbag
                .iter()
                .filter(|(attr, _)| attr.expanded() == expanded_name!(svg "values"))
            {
                let new_matrix = match operation_type {
                    OperationType::LuminanceToAlpha => unreachable!(),
                    OperationType::Matrix => {
                        let NumberList(v) =
                            NumberList::parse_str(value, NumberListLength::Exact(20)).map_err(
                                |err| {
                                    let err_str = match err {
                                        NumberListError::IncorrectNumberOfElements => {
                                            "incorrect number of elements: expected 20"
                                        }
                                        NumberListError::Parse(ref err) => &err,
                                    };

                                    NodeError::parse_error(attr, err_str)
                                },
                            )?;
                        let matrix = Matrix4x5::from_row_slice(&v);
                        let mut matrix = matrix.fixed_resize(0.0);
                        matrix[(4, 4)] = 1.0;
                        matrix
                    }
                    OperationType::Saturate => {
                        let s = parsers::number(value).attribute(attr.clone())?;
                        if s < 0.0 || s > 1.0 {
                            return Err(NodeError::value_error(attr, "expected value from 0 to 1"));
                        }

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        Matrix5::new(
                            0.213 + 0.787 * s, 0.715 - 0.715 * s, 0.072 - 0.072 * s, 0.0, 0.0,
                            0.213 - 0.213 * s, 0.715 + 0.285 * s, 0.072 - 0.072 * s, 0.0, 0.0,
                            0.213 - 0.213 * s, 0.715 - 0.715 * s, 0.072 + 0.928 * s, 0.0, 0.0,
                            0.0,               0.0,               0.0,               1.0, 0.0,
                            0.0,               0.0,               0.0,               0.0, 1.0,
                        )
                    }
                    OperationType::HueRotate => {
                        let degrees = parsers::number(value).attribute(attr.clone())?;
                        let (sin, cos) = degrees.to_radians().sin_cos();

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        let a = Matrix3::new(
                            0.213, 0.715, 0.072,
                            0.213, 0.715, 0.072,
                            0.213, 0.715, 0.072,
                        );

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        let b = Matrix3::new(
                             0.787, -0.715, -0.072,
                            -0.213,  0.285, -0.072,
                            -0.213, -0.715,  0.928,
                        );

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        let c = Matrix3::new(
                            -0.213, -0.715,  0.928,
                             0.143,  0.140, -0.283,
                            -0.787,  0.715,  0.072,
                        );

                        let top_left = a + b * cos + c * sin;

                        let mut matrix = top_left.fixed_resize(0.0);
                        matrix[(3, 3)] = 1.0;
                        matrix[(4, 4)] = 1.0;
                        matrix
                    }
                };

                self.matrix = new_matrix;
            }
        }

        Ok(())
    }
}

impl FilterEffect for FeColorMatrix {
    fn render(
        &self,
        _node: &RsvgNode,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError> {
        let input = self.base.get_input(ctx, draw_ctx)?;
        let bounds = self
            .base
            .get_bounds(ctx)
            .add_input(&input)
            .into_irect(draw_ctx);

        let mut output_surface = ImageSurface::create(
            cairo::Format::ARgb32,
            ctx.source_graphic().width(),
            ctx.source_graphic().height(),
        )?;

        let output_stride = output_surface.get_stride() as usize;
        {
            let mut output_data = output_surface.get_data().unwrap();

            for (x, y, pixel) in Pixels::new(input.surface(), bounds) {
                let alpha = f64::from(pixel.a) / 255f64;

                let pixel_vec = if alpha == 0.0 {
                    Vector5::new(0.0, 0.0, 0.0, 0.0, 1.0)
                } else {
                    Vector5::new(
                        f64::from(pixel.r) / 255f64 / alpha,
                        f64::from(pixel.g) / 255f64 / alpha,
                        f64::from(pixel.b) / 255f64 / alpha,
                        alpha,
                        1.0,
                    )
                };
                let mut new_pixel_vec = Vector5::zeros();
                self.matrix.mul_to(&pixel_vec, &mut new_pixel_vec);

                let new_alpha = clamp(new_pixel_vec[3], 0.0, 1.0);

                let premultiply = |x: f64| ((clamp(x, 0.0, 1.0) * new_alpha * 255f64) + 0.5) as u8;

                let output_pixel = Pixel {
                    r: premultiply(new_pixel_vec[0]),
                    g: premultiply(new_pixel_vec[1]),
                    b: premultiply(new_pixel_vec[2]),
                    a: ((new_alpha * 255f64) + 0.5) as u8,
                };

                output_data.set_pixel(output_stride, output_pixel, x, y);
            }
        }

        Ok(FilterResult {
            name: self.base.result.clone(),
            output: FilterOutput {
                surface: SharedImageSurface::new(output_surface, input.surface().surface_type())?,
                bounds,
            },
        })
    }

    #[inline]
    fn is_affected_by_color_interpolation_filters(&self) -> bool {
        true
    }
}

impl OperationType {
    fn parse(attr: QualName, s: &str) -> Result<Self, NodeError> {
        match s {
            "matrix" => Ok(OperationType::Matrix),
            "saturate" => Ok(OperationType::Saturate),
            "hueRotate" => Ok(OperationType::HueRotate),
            "luminanceToAlpha" => Ok(OperationType::LuminanceToAlpha),
            _ => Err(NodeError::parse_error(attr, "invalid value")),
        }
    }
}
<commit_msg>color_matrix.rs: Don't create NodeError directly<commit_after>use cairo::{self, ImageSurface};
use markup5ever::{expanded_name, local_name, namespace_url, ns, QualName};
use nalgebra::{Matrix3, Matrix4x5, Matrix5, Vector5};

use crate::drawing_ctx::DrawingCtx;
use crate::error::*;
use crate::node::{NodeResult, NodeTrait, RsvgNode};
use crate::number_list::{NumberList, NumberListError, NumberListLength};
use crate::parsers;
use crate::property_bag::PropertyBag;
use crate::surface_utils::{
    iterators::Pixels, shared_surface::SharedImageSurface, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;

use super::context::{FilterContext, FilterOutput, FilterResult};
use super::{FilterEffect, FilterError, PrimitiveWithInput};

/// Color matrix operation types.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum OperationType {
    Matrix,
    Saturate,
    HueRotate,
    LuminanceToAlpha,
}

/// The `feColorMatrix` filter primitive.
pub struct FeColorMatrix {
    base: PrimitiveWithInput,
    matrix: Matrix5<f64>,
}

impl Default for FeColorMatrix {
    /// Constructs a new `ColorMatrix` with empty properties.
    #[inline]
    fn default() -> FeColorMatrix {
        FeColorMatrix {
            base: PrimitiveWithInput::new::<Self>(),
            matrix: Matrix5::identity(),
        }
    }
}

impl NodeTrait for FeColorMatrix {
    impl_node_as_filter_effect!();

    fn set_atts(&mut self, parent: Option<&RsvgNode>, pbag: &PropertyBag<'_>) -> NodeResult {
        self.base.set_atts(parent, pbag)?;

        // First, determine the operation type.
        let mut operation_type = OperationType::Matrix;
        for (attr, value) in pbag
            .iter()
            .filter(|(attr, _)| attr.expanded() == expanded_name!(svg "type"))
        {
            operation_type = OperationType::parse(attr, value)?;
        }

        // Now read the matrix correspondingly.
        // LuminanceToAlpha doesn't accept any matrix.
        if operation_type == OperationType::LuminanceToAlpha {
            self.matrix = {
                #[cfg_attr(rustfmt, rustfmt_skip)]
                Matrix5::new(
                    0.0,    0.0,    0.0,    0.0, 0.0,
                    0.0,    0.0,    0.0,    0.0, 0.0,
                    0.0,    0.0,    0.0,    0.0, 0.0,
                    0.2125, 0.7154, 0.0721, 0.0, 0.0,
                    0.0,    0.0,    0.0,    0.0, 1.0,
                )
            };
        } else {
            for (attr, value) in pbag
                .iter()
                .filter(|(attr, _)| attr.expanded() == expanded_name!(svg "values"))
            {
                let new_matrix = match operation_type {
                    OperationType::LuminanceToAlpha => unreachable!(),
                    OperationType::Matrix => {
                        let NumberList(v) =
                            NumberList::parse_str(value, NumberListLength::Exact(20))
                                .map_err(|err| {
                                    let err_str = match err {
                                        NumberListError::IncorrectNumberOfElements => {
                                            "incorrect number of elements: expected 20"
                                        }
                                        NumberListError::Parse(ref err) => &err,
                                    };

                                    ValueErrorKind::parse_error(err_str)
                                })
                                .attribute(attr)?;
                        let matrix = Matrix4x5::from_row_slice(&v);
                        let mut matrix = matrix.fixed_resize(0.0);
                        matrix[(4, 4)] = 1.0;
                        matrix
                    }
                    OperationType::Saturate => {
                        let s = parsers::number(value).attribute(attr.clone())?;
                        if s < 0.0 || s > 1.0 {
                            return Err(ValueErrorKind::value_error("expected value from 0 to 1"))
                                .attribute(attr);
                        }

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        Matrix5::new(
                            0.213 + 0.787 * s, 0.715 - 0.715 * s, 0.072 - 0.072 * s, 0.0, 0.0,
                            0.213 - 0.213 * s, 0.715 + 0.285 * s, 0.072 - 0.072 * s, 0.0, 0.0,
                            0.213 - 0.213 * s, 0.715 - 0.715 * s, 0.072 + 0.928 * s, 0.0, 0.0,
                            0.0,               0.0,               0.0,               1.0, 0.0,
                            0.0,               0.0,               0.0,               0.0, 1.0,
                        )
                    }
                    OperationType::HueRotate => {
                        let degrees = parsers::number(value).attribute(attr.clone())?;
                        let (sin, cos) = degrees.to_radians().sin_cos();

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        let a = Matrix3::new(
                            0.213, 0.715, 0.072,
                            0.213, 0.715, 0.072,
                            0.213, 0.715, 0.072,
                        );

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        let b = Matrix3::new(
                             0.787, -0.715, -0.072,
                            -0.213,  0.285, -0.072,
                            -0.213, -0.715,  0.928,
                        );

                        #[cfg_attr(rustfmt, rustfmt_skip)]
                        let c = Matrix3::new(
                            -0.213, -0.715,  0.928,
                             0.143,  0.140, -0.283,
                            -0.787,  0.715,  0.072,
                        );

                        let top_left = a + b * cos + c * sin;

                        let mut matrix = top_left.fixed_resize(0.0);
                        matrix[(3, 3)] = 1.0;
                        matrix[(4, 4)] = 1.0;
                        matrix
                    }
                };

                self.matrix = new_matrix;
            }
        }

        Ok(())
    }
}

impl FilterEffect for FeColorMatrix {
    fn render(
        &self,
        _node: &RsvgNode,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError> {
        let input = self.base.get_input(ctx, draw_ctx)?;
        let bounds = self
            .base
            .get_bounds(ctx)
            .add_input(&input)
            .into_irect(draw_ctx);

        let mut output_surface = ImageSurface::create(
            cairo::Format::ARgb32,
            ctx.source_graphic().width(),
            ctx.source_graphic().height(),
        )?;

        let output_stride = output_surface.get_stride() as usize;
        {
            let mut output_data = output_surface.get_data().unwrap();

            for (x, y, pixel) in Pixels::new(input.surface(), bounds) {
                let alpha = f64::from(pixel.a) / 255f64;

                let pixel_vec = if alpha == 0.0 {
                    Vector5::new(0.0, 0.0, 0.0, 0.0, 1.0)
                } else {
                    Vector5::new(
                        f64::from(pixel.r) / 255f64 / alpha,
                        f64::from(pixel.g) / 255f64 / alpha,
                        f64::from(pixel.b) / 255f64 / alpha,
                        alpha,
                        1.0,
                    )
                };
                let mut new_pixel_vec = Vector5::zeros();
                self.matrix.mul_to(&pixel_vec, &mut new_pixel_vec);

                let new_alpha = clamp(new_pixel_vec[3], 0.0, 1.0);

                let premultiply = |x: f64| ((clamp(x, 0.0, 1.0) * new_alpha * 255f64) + 0.5) as u8;

                let output_pixel = Pixel {
                    r: premultiply(new_pixel_vec[0]),
                    g: premultiply(new_pixel_vec[1]),
                    b: premultiply(new_pixel_vec[2]),
                    a: ((new_alpha * 255f64) + 0.5) as u8,
                };

                output_data.set_pixel(output_stride, output_pixel, x, y);
            }
        }

        Ok(FilterResult {
            name: self.base.result.clone(),
            output: FilterOutput {
                surface: SharedImageSurface::new(output_surface, input.surface().surface_type())?,
                bounds,
            },
        })
    }

    #[inline]
    fn is_affected_by_color_interpolation_filters(&self) -> bool {
        true
    }
}

impl OperationType {
    fn parse(attr: QualName, s: &str) -> Result<Self, NodeError> {
        match s {
            "matrix" => Ok(OperationType::Matrix),
            "saturate" => Ok(OperationType::Saturate),
            "hueRotate" => Ok(OperationType::HueRotate),
            "luminanceToAlpha" => Ok(OperationType::LuminanceToAlpha),
            _ => Err(ValueErrorKind::parse_error("invalid value")).attribute(attr),
        }
    }
}
<|endoftext|>


<commit_before>//! Logistic Regression module
//!
//! Contains implemention of logistic regression using
//! gradient descent optimization.
//!
//! The regressor will automatically add the intercept term
//! so you do not need to format the input matrices yourself.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::learning::logistic_reg::LogisticRegressor;
//! use rusty_machine::learning::SupModel;
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::linalg::vector::Vector;
//!
//! let inputs = Matrix::new(4,1,vec![1.0,3.0,5.0,7.0]);
//! let targets = Vector::new(vec![0.,0.,1.,1.]);
//!
//! let mut log_mod = LogisticRegressor::default();
//!
//! // Train the model
//! log_mod.train(&inputs, &targets);
//!
//! // Now we'll predict a new point
//! let new_point = Matrix::new(1,1,vec![10.]);
//! let output = log_mod.predict(&new_point);
//!
//! // Hopefully we classified our new point correctly!
//! assert!(output[0] > 0.5, "Our classifier isn't very good!");
//! ```
//!
//! We could have been more specific about the learning of the model
//! by using the `new` constructor instead. This allows us to provide
//! a `GradientDesc` object with custom parameters.

use learning::SupModel;
use linalg::matrix::Matrix;
use linalg::vector::Vector;
use learning::toolkit::activ_fn::ActivationFunc;
use learning::toolkit::activ_fn::Sigmoid;
use learning::toolkit::cost_fn::CostFunc;
use learning::toolkit::cost_fn::CrossEntropyError;
use learning::optim::grad_desc::GradientDesc;
use learning::optim::OptimAlgorithm;
use learning::optim::Optimizable;

/// Logistic Regression Model.
///
/// Contains option for optimized parameter.
pub struct LogisticRegressor {
    /// The parameters for the regression model.
    parameters: Option<Vector<f64>>,
    gd: GradientDesc,
}

impl Default for LogisticRegressor {
    fn default() -> LogisticRegressor {
        LogisticRegressor {
            parameters: None,
            gd: GradientDesc::default(),
        }
    }
}
impl LogisticRegressor {
    /// Constructs untrained logistic regression model.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::logistic_reg::LogisticRegressor;
    /// use rusty_machine::learning::optim::grad_desc::GradientDesc;
    ///
    /// let gd = GradientDesc::default();
    /// let mut logistic_mod = LogisticRegressor::new(gd);
    /// ```
    pub fn new(gd: GradientDesc) -> LogisticRegressor {
        LogisticRegressor {
            parameters: None,
            gd: gd,
        }
    }

    /// Get the parameters from the model.
    ///
    /// Returns an option that is None if the model has not been trained.
    pub fn parameters(&self) -> Option<Vector<f64>> {
        match self.parameters {
            None => None,
            Some(ref x) => Some(x.clone()),
        }
    }
}

impl SupModel<Matrix<f64>, Vector<f64>> for LogisticRegressor {
    /// Train the logistic regression model.
    ///
    /// Takes training data and output values as input.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::logistic_reg::LogisticRegressor;
    /// use rusty_machine::linalg::matrix::Matrix;
    /// use rusty_machine::linalg::vector::Vector;
    /// use rusty_machine::learning::SupModel;
    ///
    /// let mut logistic_mod = LogisticRegressor::default();
    /// let inputs = Matrix::new(3,2, vec![1.0, 2.0, 1.0, 3.0, 1.0, 4.0]);
    /// let targets = Vector::new(vec![5.0, 6.0, 7.0]);
    ///
    /// logistic_mod.train(&inputs, &targets);
    /// ```
    fn train(&mut self, inputs: &Matrix<f64>, targets: &Vector<f64>) {
        let ones = Matrix::<f64>::ones(inputs.rows(), 1);
        let full_inputs = ones.hcat(inputs);

        let initial_params = vec![0.5; full_inputs.cols()];

        let optimal_w = self.gd.optimize(self, &initial_params[..], &full_inputs, targets);
        self.parameters = Some(Vector::new(optimal_w));
    }

    /// Predict output value from input data.
    ///
    /// Model must be trained before prediction can be made.
    fn predict(&self, inputs: &Matrix<f64>) -> Vector<f64> {
        if let Some(ref v) = self.parameters {
            let ones = Matrix::<f64>::ones(inputs.rows(), 1);
            let full_inputs = ones.hcat(inputs);
            (full_inputs * v).apply(&Sigmoid::func)
        } else {
            panic!("Model has not been trained.");
        }
    }
}

impl Optimizable for LogisticRegressor {
    type Inputs = Matrix<f64>;
    type Targets = Vector<f64>;

    fn compute_grad(&self,
                    params: &[f64],
                    inputs: &Matrix<f64>,
                    targets: &Vector<f64>)
                    -> (f64, Vec<f64>) {

        let beta_vec = Vector::new(params.to_vec());
        let outputs = (inputs * beta_vec).apply(&Sigmoid::func);

        let cost = CrossEntropyError::cost(&outputs, targets);
        let grad = (inputs.transpose() * (outputs - targets)) / (inputs.rows() as f64);

        println!("Cost is {0}", cost);

        (cost, grad.into_vec())
    }
}
<commit_msg>Cleaning up printing<commit_after>//! Logistic Regression module
//!
//! Contains implemention of logistic regression using
//! gradient descent optimization.
//!
//! The regressor will automatically add the intercept term
//! so you do not need to format the input matrices yourself.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::learning::logistic_reg::LogisticRegressor;
//! use rusty_machine::learning::SupModel;
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::linalg::vector::Vector;
//!
//! let inputs = Matrix::new(4,1,vec![1.0,3.0,5.0,7.0]);
//! let targets = Vector::new(vec![0.,0.,1.,1.]);
//!
//! let mut log_mod = LogisticRegressor::default();
//!
//! // Train the model
//! log_mod.train(&inputs, &targets);
//!
//! // Now we'll predict a new point
//! let new_point = Matrix::new(1,1,vec![10.]);
//! let output = log_mod.predict(&new_point);
//!
//! // Hopefully we classified our new point correctly!
//! assert!(output[0] > 0.5, "Our classifier isn't very good!");
//! ```
//!
//! We could have been more specific about the learning of the model
//! by using the `new` constructor instead. This allows us to provide
//! a `GradientDesc` object with custom parameters.

use learning::SupModel;
use linalg::matrix::Matrix;
use linalg::vector::Vector;
use learning::toolkit::activ_fn::ActivationFunc;
use learning::toolkit::activ_fn::Sigmoid;
use learning::toolkit::cost_fn::CostFunc;
use learning::toolkit::cost_fn::CrossEntropyError;
use learning::optim::grad_desc::GradientDesc;
use learning::optim::OptimAlgorithm;
use learning::optim::Optimizable;

/// Logistic Regression Model.
///
/// Contains option for optimized parameter.
pub struct LogisticRegressor {
    /// The parameters for the regression model.
    parameters: Option<Vector<f64>>,
    gd: GradientDesc,
}

impl Default for LogisticRegressor {
    fn default() -> LogisticRegressor {
        LogisticRegressor {
            parameters: None,
            gd: GradientDesc::default(),
        }
    }
}
impl LogisticRegressor {
    /// Constructs untrained logistic regression model.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::logistic_reg::LogisticRegressor;
    /// use rusty_machine::learning::optim::grad_desc::GradientDesc;
    ///
    /// let gd = GradientDesc::default();
    /// let mut logistic_mod = LogisticRegressor::new(gd);
    /// ```
    pub fn new(gd: GradientDesc) -> LogisticRegressor {
        LogisticRegressor {
            parameters: None,
            gd: gd,
        }
    }

    /// Get the parameters from the model.
    ///
    /// Returns an option that is None if the model has not been trained.
    pub fn parameters(&self) -> Option<Vector<f64>> {
        match self.parameters {
            None => None,
            Some(ref x) => Some(x.clone()),
        }
    }
}

impl SupModel<Matrix<f64>, Vector<f64>> for LogisticRegressor {
    /// Train the logistic regression model.
    ///
    /// Takes training data and output values as input.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::logistic_reg::LogisticRegressor;
    /// use rusty_machine::linalg::matrix::Matrix;
    /// use rusty_machine::linalg::vector::Vector;
    /// use rusty_machine::learning::SupModel;
    ///
    /// let mut logistic_mod = LogisticRegressor::default();
    /// let inputs = Matrix::new(3,2, vec![1.0, 2.0, 1.0, 3.0, 1.0, 4.0]);
    /// let targets = Vector::new(vec![5.0, 6.0, 7.0]);
    ///
    /// logistic_mod.train(&inputs, &targets);
    /// ```
    fn train(&mut self, inputs: &Matrix<f64>, targets: &Vector<f64>) {
        let ones = Matrix::<f64>::ones(inputs.rows(), 1);
        let full_inputs = ones.hcat(inputs);

        let initial_params = vec![0.5; full_inputs.cols()];

        let optimal_w = self.gd.optimize(self, &initial_params[..], &full_inputs, targets);
        self.parameters = Some(Vector::new(optimal_w));
    }

    /// Predict output value from input data.
    ///
    /// Model must be trained before prediction can be made.
    fn predict(&self, inputs: &Matrix<f64>) -> Vector<f64> {
        if let Some(ref v) = self.parameters {
            let ones = Matrix::<f64>::ones(inputs.rows(), 1);
            let full_inputs = ones.hcat(inputs);
            (full_inputs * v).apply(&Sigmoid::func)
        } else {
            panic!("Model has not been trained.");
        }
    }
}

impl Optimizable for LogisticRegressor {
    type Inputs = Matrix<f64>;
    type Targets = Vector<f64>;

    fn compute_grad(&self,
                    params: &[f64],
                    inputs: &Matrix<f64>,
                    targets: &Vector<f64>)
                    -> (f64, Vec<f64>) {

        let beta_vec = Vector::new(params.to_vec());
        let outputs = (inputs * beta_vec).apply(&Sigmoid::func);

        let cost = CrossEntropyError::cost(&outputs, targets);
        let grad = (inputs.transpose() * (outputs - targets)) / (inputs.rows() as f64);

        (cost, grad.into_vec())
    }
}
<|endoftext|>


<commit_before><commit_msg>CHD: Invalidate location when setting one failed<commit_after><|endoftext|>


<commit_before><commit_msg>forgot to call format<commit_after><|endoftext|>


<commit_before><commit_msg>document some things<commit_after><|endoftext|>


<commit_before><commit_msg>Fix for latest nightly<commit_after><|endoftext|>


<commit_before><commit_msg>Errors are strings now.<commit_after><|endoftext|>


<commit_before><commit_msg>feat: FswSession::{builder, builder_params}<commit_after><|endoftext|>


<commit_before><commit_msg>Added destruction_queue<commit_after><|endoftext|>


<commit_before><commit_msg>Fix build break<commit_after><|endoftext|>


<commit_before><commit_msg>Create function that returns handler errors<commit_after><|endoftext|>


<commit_before><commit_msg>Remove useless test in lib.rs<commit_after><|endoftext|>


<commit_before><commit_msg>compat<commit_after><|endoftext|>


<commit_before><commit_msg>Set::is_disjoint<commit_after><|endoftext|>


<commit_before>//! Complete commands within shells

/// Complete commands within bash
pub mod bash {
    use std::ffi::OsString;
    use std::io::Write;

    use unicode_xid::UnicodeXID;

    #[derive(clap::Subcommand)]
    #[clap(hide = true)]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub enum CompleteCommand {
        /// Register shell completions for this program
        Complete(CompleteArgs),
    }

    #[derive(clap::Args)]
    #[clap(group = clap::ArgGroup::new("complete").multiple(true).conflicts_with("register"))]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub struct CompleteArgs {
        /// Path to write completion-registration to
        #[clap(long, required = true, parse(from_os_str))]
        register: Option<std::path::PathBuf>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        index: Option<usize>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        ifs: Option<String>,

        #[clap(
            long = "type",
            required = true,
            arg_enum,
            hide_short_help = true,
            group = "complete"
        )]
        comp_type: Option<CompType>,

        #[clap(long, hide_short_help = true, group = "complete")]
        space: bool,

        #[clap(
            long,
            conflicts_with = "space",
            hide_short_help = true,
            group = "complete"
        )]
        no_space: bool,

        #[clap(raw = true, hide_short_help = true, group = "complete")]
        comp_words: Vec<OsString>,
    }

    impl CompleteCommand {
        /// Process the completion request
        pub fn complete(&self, cmd: &mut clap::Command) -> std::convert::Infallible {
            self.try_complete(cmd).unwrap_or_else(|e| e.exit());
            std::process::exit(0)
        }

        /// Process the completion request
        pub fn try_complete(&self, cmd: &mut clap::Command) -> clap::Result<()> {
            debug!("CompleteCommand::try_complete: {:?}", self);
            let CompleteCommand::Complete(args) = self;
            if let Some(out_path) = args.register.as_deref() {
                let mut buf = Vec::new();
                let name = cmd.get_name();
                let bin = cmd.get_bin_name().unwrap_or(cmd.get_name());
                register(name, [bin], bin, &Behavior::default(), &mut buf)?;
                if out_path == std::path::Path::new("-") {
                    std::io::stdout().write(&buf)?;
                } else {
                    if out_path.is_dir() {
                        let out_path = out_path.join(file_name(name));
                        std::fs::write(out_path, buf)?;
                    } else {
                        std::fs::write(out_path, buf)?;
                    }
                }
            } else {
                let index = args.index.unwrap_or_default();
                let comp_type = args.comp_type.unwrap_or_default();
                let space = match (args.space, args.no_space) {
                    (true, false) => Some(true),
                    (false, true) => Some(false),
                    (true, true) => {
                        unreachable!("`--space` and `--no-space` set, clap should prevent this")
                    }
                    (false, false) => None,
                }
                .unwrap();
                let current_dir = std::env::current_dir().ok();
                let completions = complete(
                    cmd,
                    args.comp_words.clone(),
                    index,
                    comp_type,
                    space,
                    current_dir.as_deref(),
                )?;

                let mut buf = Vec::new();
                for completion in &completions {
                    write!(&mut buf, "{}", completion.to_string_lossy())?;
                    write!(&mut buf, "{}", args.ifs.as_deref().unwrap_or("\n"))?;
                }
                std::io::stdout().write(&buf)?;
            }

            Ok(())
        }
    }

    /// The recommended file name for the registration code
    pub fn file_name(name: &str) -> String {
        format!("{}.bash", name)
    }

    /// Define the completion behavior
    pub enum Behavior {
        /// Bare bones behavior
        Minimal,
        /// Fallback to readline behavior when no matches are generated
        Readline,
        /// Customize bash's completion behavior
        Custom(String),
    }

    impl Default for Behavior {
        fn default() -> Self {
            Self::Readline
        }
    }

    /// Generate code to register the dynamic completion
    pub fn register(
        name: &str,
        executables: impl IntoIterator<Item = impl AsRef<str>>,
        completer: &str,
        behavior: &Behavior,
        buf: &mut dyn Write,
    ) -> Result<(), std::io::Error> {
        let escaped_name = name.replace("-", "_");
        debug_assert!(
            escaped_name.chars().all(|c| c.is_xid_continue()),
            "`name` must be an identifier, got `{}`",
            escaped_name
        );
        let mut upper_name = escaped_name.clone();
        upper_name.make_ascii_uppercase();

        let executables = executables
            .into_iter()
            .map(|s| shlex::quote(s.as_ref()).into_owned())
            .collect::<Vec<_>>()
            .join(" ");

        let options = match behavior {
            Behavior::Minimal => "-o nospace -o bashdefault",
            Behavior::Readline => "-o nospace -o default -o bashdefault",
            Behavior::Custom(c) => c.as_str(),
        };

        let completer = shlex::quote(completer);

        let script = r#"
_clap_complete_NAME() {
    local IFS=$'\013'
    local SUPPRESS_SPACE=0
    if compopt +o nospace 2> /dev/null; then
        SUPPRESS_SPACE=1
    fi
    if [[ ${SUPPRESS_SPACE} == 1 ]]; then
        SPACE_ARG="--no-space"
    else
        SPACE_ARG="--space"
    fi
    COMPREPLY=( $("COMPLETER" complete --index ${COMP_CWORD} --type ${COMP_TYPE} ${SPACE_ARG} --ifs="$IFS" -- "${COMP_WORDS[@]}") )
    if [[ $? != 0 ]]; then
        unset COMPREPLY
    elif [[ $SUPPRESS_SPACE == 1 ]] && [[ "${COMPREPLY-}" =~ [=/:]$ ]]; then
        compopt -o nospace
    fi
}
complete OPTIONS -F _clap_complete_NAME EXECUTABLES
"#
        .replace("NAME", &escaped_name)
        .replace("EXECUTABLES", &executables)
        .replace("OPTIONS", options)
        .replace("COMPLETER", &completer)
        .replace("UPPER", &upper_name);

        writeln!(buf, "{}", script)?;
        Ok(())
    }

    /// Type of completion attempted that caused a completion function to be called
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    #[non_exhaustive]
    pub enum CompType {
        /// Normal completion
        Normal,
        /// List completions after successive tabs
        Successive,
        /// List alternatives on partial word completion
        Alternatives,
        /// List completions if the word is not unmodified
        Unmodified,
        /// Menu completion
        Menu,
    }

    impl clap::ArgEnum for CompType {
        fn value_variants<'a>() -> &'a [Self] {
            &[
                Self::Normal,
                Self::Successive,
                Self::Alternatives,
                Self::Unmodified,
                Self::Menu,
            ]
        }
        fn to_possible_value<'a>(&self) -> ::std::option::Option<clap::PossibleValue<'a>> {
            match self {
                Self::Normal => {
                    let value = "9";
                    debug_assert_eq!(b'\t'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("normal")
                            .help("Normal completion"),
                    )
                }
                Self::Successive => {
                    let value = "63";
                    debug_assert_eq!(b'?'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("successive")
                            .help("List completions after successive tabs"),
                    )
                }
                Self::Alternatives => {
                    let value = "33";
                    debug_assert_eq!(b'!'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("alternatives")
                            .help("List alternatives on partial word completion"),
                    )
                }
                Self::Unmodified => {
                    let value = "64";
                    debug_assert_eq!(b'@'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("unmodified")
                            .help("List completions if the word is not unmodified"),
                    )
                }
                Self::Menu => {
                    let value = "37";
                    debug_assert_eq!(b'%'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("menu")
                            .help("Menu completion"),
                    )
                }
            }
        }
    }

    impl Default for CompType {
        fn default() -> Self {
            Self::Normal
        }
    }

    /// Complete the command specified
    pub fn complete(
        cmd: &mut clap::Command,
        args: Vec<std::ffi::OsString>,
        arg_index: usize,
        _comp_type: CompType,
        _trailing_space: bool,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        cmd.build();

        let raw_args = clap_lex::RawArgs::new(args.into_iter());
        let mut cursor = raw_args.cursor();
        let mut target_cursor = raw_args.cursor();
        raw_args.seek(
            &mut target_cursor,
            clap_lex::SeekFrom::Start(arg_index as u64),
        );
        // As we loop, `cursor` will always be pointing to the next item
        raw_args.next_os(&mut target_cursor);

        // TODO: Multicall support
        if !cmd.is_no_binary_name_set() {
            raw_args.next_os(&mut cursor);
        }

        let mut current_cmd = &*cmd;
        let mut pos_index = 0;
        while let Some(arg) = raw_args.next(&mut cursor) {
            if cursor == target_cursor {
                return complete_new(current_cmd, current_dir, pos_index);
            }
            if let Ok(value) = arg.to_value() {
                if let Some(next_cmd) = current_cmd.find_subcommand(value) {
                    current_cmd = next_cmd;
                    pos_index = 0;
                    continue;
                }
            }

            pos_index += 1;
        }

        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "No completion generated",
        ))
    }

    fn complete_new(
        cmd: &clap::Command,
        current_dir: Option<&std::path::Path>,
        pos_index: usize,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        let mut completions = Vec::new();

        completions.extend(
            crate::generator::utils::longs_and_visible_aliases(cmd)
                .into_iter()
                .map(|f| format!("--{}", f).into()),
        );
        completions.extend(
            crate::generator::utils::shorts_and_visible_aliases(cmd)
                .into_iter()
                .map(|f| format!("-{}", f).into()),
        );

        let mut positionals = Vec::new();
        if let Some(positional) = cmd
            .get_positionals()
            .find(|p| p.get_index() == Some(pos_index))
        {
            if let Some(possible_values) = positional.get_possible_values() {
                positionals.extend(possible_values.into_iter().map(|p| p.get_name().into()));
            } else {
                match positional.get_value_hint() {
                    clap::ValueHint::Other => {
                        // Should not complete
                    }
                    clap::ValueHint::Unknown | clap::ValueHint::AnyPath => {
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                    clap::ValueHint::FilePath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_file()));
                    }
                    clap::ValueHint::DirPath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_dir()));
                    }
                    clap::ValueHint::ExecutablePath => {
                        use is_executable::IsExecutable;
                        positionals.extend(complete_path(current_dir, |p| p.is_executable()));
                    }
                    clap::ValueHint::CommandName
                    | clap::ValueHint::CommandString
                    | clap::ValueHint::CommandWithArguments
                    | clap::ValueHint::Username
                    | clap::ValueHint::Hostname
                    | clap::ValueHint::Url
                    | clap::ValueHint::EmailAddress => {
                        // No completion implementation
                    }
                    _ => {
                        // Safe-ish fallback
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                }
                positionals.sort();
            }
            completions.extend(positionals);
        }

        completions.extend(all_subcommands(cmd).into_iter().map(|s| s));

        Ok(completions)
    }

    fn complete_path(
        current_dir: Option<&std::path::Path>,
        is_wanted: impl Fn(&std::path::Path) -> bool,
    ) -> Vec<OsString> {
        let mut completions = Vec::new();

        let current_dir = match current_dir {
            Some(current_dir) => current_dir,
            None => {
                return Vec::new();
            }
        };
        for entry in std::fs::read_dir(current_dir)
            .ok()
            .into_iter()
            .flatten()
            .filter_map(Result::ok)
        {
            let mut path = entry.path();
            if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
                path.push(""); // Ensure trailing `/`
                let suggestion = path.strip_prefix(current_dir).unwrap();
                completions.push(suggestion.as_os_str().to_owned());
            } else {
                if is_wanted(&path) {
                    let suggestion = path.strip_prefix(current_dir).unwrap();
                    completions.push(suggestion.as_os_str().to_owned());
                }
            }
        }

        completions
    }

    fn all_subcommands(cmd: &clap::Command) -> Vec<OsString> {
        debug!("all_subcommands");

        let mut scs = crate::generator::utils::all_subcommands(cmd)
            .iter()
            .map(|x| OsString::from(&x.0))
            .collect::<Vec<_>>();
        scs.sort();
        scs.dedup();
        scs
    }
}
<commit_msg>feat(complete): Escaped args are always positional<commit_after>//! Complete commands within shells

/// Complete commands within bash
pub mod bash {
    use std::ffi::OsString;
    use std::io::Write;

    use unicode_xid::UnicodeXID;

    #[derive(clap::Subcommand)]
    #[clap(hide = true)]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub enum CompleteCommand {
        /// Register shell completions for this program
        Complete(CompleteArgs),
    }

    #[derive(clap::Args)]
    #[clap(group = clap::ArgGroup::new("complete").multiple(true).conflicts_with("register"))]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub struct CompleteArgs {
        /// Path to write completion-registration to
        #[clap(long, required = true, parse(from_os_str))]
        register: Option<std::path::PathBuf>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        index: Option<usize>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        ifs: Option<String>,

        #[clap(
            long = "type",
            required = true,
            arg_enum,
            hide_short_help = true,
            group = "complete"
        )]
        comp_type: Option<CompType>,

        #[clap(long, hide_short_help = true, group = "complete")]
        space: bool,

        #[clap(
            long,
            conflicts_with = "space",
            hide_short_help = true,
            group = "complete"
        )]
        no_space: bool,

        #[clap(raw = true, hide_short_help = true, group = "complete")]
        comp_words: Vec<OsString>,
    }

    impl CompleteCommand {
        /// Process the completion request
        pub fn complete(&self, cmd: &mut clap::Command) -> std::convert::Infallible {
            self.try_complete(cmd).unwrap_or_else(|e| e.exit());
            std::process::exit(0)
        }

        /// Process the completion request
        pub fn try_complete(&self, cmd: &mut clap::Command) -> clap::Result<()> {
            debug!("CompleteCommand::try_complete: {:?}", self);
            let CompleteCommand::Complete(args) = self;
            if let Some(out_path) = args.register.as_deref() {
                let mut buf = Vec::new();
                let name = cmd.get_name();
                let bin = cmd.get_bin_name().unwrap_or(cmd.get_name());
                register(name, [bin], bin, &Behavior::default(), &mut buf)?;
                if out_path == std::path::Path::new("-") {
                    std::io::stdout().write(&buf)?;
                } else {
                    if out_path.is_dir() {
                        let out_path = out_path.join(file_name(name));
                        std::fs::write(out_path, buf)?;
                    } else {
                        std::fs::write(out_path, buf)?;
                    }
                }
            } else {
                let index = args.index.unwrap_or_default();
                let comp_type = args.comp_type.unwrap_or_default();
                let space = match (args.space, args.no_space) {
                    (true, false) => Some(true),
                    (false, true) => Some(false),
                    (true, true) => {
                        unreachable!("`--space` and `--no-space` set, clap should prevent this")
                    }
                    (false, false) => None,
                }
                .unwrap();
                let current_dir = std::env::current_dir().ok();
                let completions = complete(
                    cmd,
                    args.comp_words.clone(),
                    index,
                    comp_type,
                    space,
                    current_dir.as_deref(),
                )?;

                let mut buf = Vec::new();
                for completion in &completions {
                    write!(&mut buf, "{}", completion.to_string_lossy())?;
                    write!(&mut buf, "{}", args.ifs.as_deref().unwrap_or("\n"))?;
                }
                std::io::stdout().write(&buf)?;
            }

            Ok(())
        }
    }

    /// The recommended file name for the registration code
    pub fn file_name(name: &str) -> String {
        format!("{}.bash", name)
    }

    /// Define the completion behavior
    pub enum Behavior {
        /// Bare bones behavior
        Minimal,
        /// Fallback to readline behavior when no matches are generated
        Readline,
        /// Customize bash's completion behavior
        Custom(String),
    }

    impl Default for Behavior {
        fn default() -> Self {
            Self::Readline
        }
    }

    /// Generate code to register the dynamic completion
    pub fn register(
        name: &str,
        executables: impl IntoIterator<Item = impl AsRef<str>>,
        completer: &str,
        behavior: &Behavior,
        buf: &mut dyn Write,
    ) -> Result<(), std::io::Error> {
        let escaped_name = name.replace("-", "_");
        debug_assert!(
            escaped_name.chars().all(|c| c.is_xid_continue()),
            "`name` must be an identifier, got `{}`",
            escaped_name
        );
        let mut upper_name = escaped_name.clone();
        upper_name.make_ascii_uppercase();

        let executables = executables
            .into_iter()
            .map(|s| shlex::quote(s.as_ref()).into_owned())
            .collect::<Vec<_>>()
            .join(" ");

        let options = match behavior {
            Behavior::Minimal => "-o nospace -o bashdefault",
            Behavior::Readline => "-o nospace -o default -o bashdefault",
            Behavior::Custom(c) => c.as_str(),
        };

        let completer = shlex::quote(completer);

        let script = r#"
_clap_complete_NAME() {
    local IFS=$'\013'
    local SUPPRESS_SPACE=0
    if compopt +o nospace 2> /dev/null; then
        SUPPRESS_SPACE=1
    fi
    if [[ ${SUPPRESS_SPACE} == 1 ]]; then
        SPACE_ARG="--no-space"
    else
        SPACE_ARG="--space"
    fi
    COMPREPLY=( $("COMPLETER" complete --index ${COMP_CWORD} --type ${COMP_TYPE} ${SPACE_ARG} --ifs="$IFS" -- "${COMP_WORDS[@]}") )
    if [[ $? != 0 ]]; then
        unset COMPREPLY
    elif [[ $SUPPRESS_SPACE == 1 ]] && [[ "${COMPREPLY-}" =~ [=/:]$ ]]; then
        compopt -o nospace
    fi
}
complete OPTIONS -F _clap_complete_NAME EXECUTABLES
"#
        .replace("NAME", &escaped_name)
        .replace("EXECUTABLES", &executables)
        .replace("OPTIONS", options)
        .replace("COMPLETER", &completer)
        .replace("UPPER", &upper_name);

        writeln!(buf, "{}", script)?;
        Ok(())
    }

    /// Type of completion attempted that caused a completion function to be called
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    #[non_exhaustive]
    pub enum CompType {
        /// Normal completion
        Normal,
        /// List completions after successive tabs
        Successive,
        /// List alternatives on partial word completion
        Alternatives,
        /// List completions if the word is not unmodified
        Unmodified,
        /// Menu completion
        Menu,
    }

    impl clap::ArgEnum for CompType {
        fn value_variants<'a>() -> &'a [Self] {
            &[
                Self::Normal,
                Self::Successive,
                Self::Alternatives,
                Self::Unmodified,
                Self::Menu,
            ]
        }
        fn to_possible_value<'a>(&self) -> ::std::option::Option<clap::PossibleValue<'a>> {
            match self {
                Self::Normal => {
                    let value = "9";
                    debug_assert_eq!(b'\t'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("normal")
                            .help("Normal completion"),
                    )
                }
                Self::Successive => {
                    let value = "63";
                    debug_assert_eq!(b'?'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("successive")
                            .help("List completions after successive tabs"),
                    )
                }
                Self::Alternatives => {
                    let value = "33";
                    debug_assert_eq!(b'!'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("alternatives")
                            .help("List alternatives on partial word completion"),
                    )
                }
                Self::Unmodified => {
                    let value = "64";
                    debug_assert_eq!(b'@'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("unmodified")
                            .help("List completions if the word is not unmodified"),
                    )
                }
                Self::Menu => {
                    let value = "37";
                    debug_assert_eq!(b'%'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("menu")
                            .help("Menu completion"),
                    )
                }
            }
        }
    }

    impl Default for CompType {
        fn default() -> Self {
            Self::Normal
        }
    }

    /// Complete the command specified
    pub fn complete(
        cmd: &mut clap::Command,
        args: Vec<std::ffi::OsString>,
        arg_index: usize,
        _comp_type: CompType,
        _trailing_space: bool,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        cmd.build();

        let raw_args = clap_lex::RawArgs::new(args.into_iter());
        let mut cursor = raw_args.cursor();
        let mut target_cursor = raw_args.cursor();
        raw_args.seek(
            &mut target_cursor,
            clap_lex::SeekFrom::Start(arg_index as u64),
        );
        // As we loop, `cursor` will always be pointing to the next item
        raw_args.next_os(&mut target_cursor);

        // TODO: Multicall support
        if !cmd.is_no_binary_name_set() {
            raw_args.next_os(&mut cursor);
        }

        let mut current_cmd = &*cmd;
        let mut pos_index = 0;
        let mut is_escaped = false;
        while let Some(arg) = raw_args.next(&mut cursor) {
            if cursor == target_cursor {
                return complete_new(current_cmd, current_dir, pos_index, is_escaped);
            }
            if let Ok(value) = arg.to_value() {
                if let Some(next_cmd) = current_cmd.find_subcommand(value) {
                    current_cmd = next_cmd;
                    pos_index = 0;
                    continue;
                }
            }

            if arg.is_escape() {
                is_escaped = true;
            } else {
                pos_index += 1;
            }
        }

        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "No completion generated",
        ))
    }

    fn complete_new(
        cmd: &clap::Command,
        current_dir: Option<&std::path::Path>,
        pos_index: usize,
        is_escaped: bool,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        let mut completions = Vec::new();

        if !is_escaped {
            completions.extend(
                crate::generator::utils::longs_and_visible_aliases(cmd)
                    .into_iter()
                    .map(|f| format!("--{}", f).into()),
            );
            completions.extend(
                crate::generator::utils::shorts_and_visible_aliases(cmd)
                    .into_iter()
                    .map(|f| format!("-{}", f).into()),
            );
        }

        let mut positionals = Vec::new();
        if let Some(positional) = cmd
            .get_positionals()
            .find(|p| p.get_index() == Some(pos_index))
        {
            if let Some(possible_values) = positional.get_possible_values() {
                positionals.extend(possible_values.into_iter().map(|p| p.get_name().into()));
            } else {
                match positional.get_value_hint() {
                    clap::ValueHint::Other => {
                        // Should not complete
                    }
                    clap::ValueHint::Unknown | clap::ValueHint::AnyPath => {
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                    clap::ValueHint::FilePath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_file()));
                    }
                    clap::ValueHint::DirPath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_dir()));
                    }
                    clap::ValueHint::ExecutablePath => {
                        use is_executable::IsExecutable;
                        positionals.extend(complete_path(current_dir, |p| p.is_executable()));
                    }
                    clap::ValueHint::CommandName
                    | clap::ValueHint::CommandString
                    | clap::ValueHint::CommandWithArguments
                    | clap::ValueHint::Username
                    | clap::ValueHint::Hostname
                    | clap::ValueHint::Url
                    | clap::ValueHint::EmailAddress => {
                        // No completion implementation
                    }
                    _ => {
                        // Safe-ish fallback
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                }
                positionals.sort();
            }
            completions.extend(positionals);
        }

        if !is_escaped {
            completions.extend(all_subcommands(cmd).into_iter().map(|s| s));
        }

        Ok(completions)
    }

    fn complete_path(
        current_dir: Option<&std::path::Path>,
        is_wanted: impl Fn(&std::path::Path) -> bool,
    ) -> Vec<OsString> {
        let mut completions = Vec::new();

        let current_dir = match current_dir {
            Some(current_dir) => current_dir,
            None => {
                return Vec::new();
            }
        };
        for entry in std::fs::read_dir(current_dir)
            .ok()
            .into_iter()
            .flatten()
            .filter_map(Result::ok)
        {
            let mut path = entry.path();
            if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
                path.push(""); // Ensure trailing `/`
                let suggestion = path.strip_prefix(current_dir).unwrap();
                completions.push(suggestion.as_os_str().to_owned());
            } else {
                if is_wanted(&path) {
                    let suggestion = path.strip_prefix(current_dir).unwrap();
                    completions.push(suggestion.as_os_str().to_owned());
                }
            }
        }

        completions
    }

    fn all_subcommands(cmd: &clap::Command) -> Vec<OsString> {
        debug!("all_subcommands");

        let mut scs = crate::generator::utils::all_subcommands(cmd)
            .iter()
            .map(|x| OsString::from(&x.0))
            .collect::<Vec<_>>();
        scs.sort();
        scs.dedup();
        scs
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
Sampling from random distributions.

This is a generalization of `Rand` to allow parameters to control the
exact properties of the generated values, e.g. the mean and standard
deviation of a normal distribution. The `Sample` trait is the most
general, and allows for generating values that change some state
internally. The `IndependentSample` trait is for generating values
that do not need to record state.

*/

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Panics if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new<'a>(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R:Rng>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            pdf: |f64|: 'static -> f64,
            zero_case: |&mut R, f64|: 'static -> f64)
            -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

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

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new(&mut []);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<commit_msg>/*! -> //!<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Sampling from random distributions.
//!
//! This is a generalization of `Rand` to allow parameters to control the
//! exact properties of the generated values, e.g. the mean and standard
//! deviation of a normal distribution. The `Sample` trait is the most
//! general, and allows for generating values that change some state
//! internally. The `IndependentSample` trait is for generating values
//! that do not need to record state.

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Panics if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new<'a>(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R:Rng>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            pdf: |f64|: 'static -> f64,
            zero_case: |&mut R, f64|: 'static -> f64)
            -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

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

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new(&mut []);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<|endoftext|>