<commit_before>//! The parser module is responsible for parsing FITS files.

use std::str;
use std::str::FromStr;
use super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};

named!(#[doc = "Will parse data from a FITS file into a `Fits` structure"], pub fits<&[u8], Fits>,
       do_parse!(
           ph: primary_header >>
               many0!(take!(2880)) >>
               (Fits::new(ph))
       ));

named!(primary_header<&[u8], PrimaryHeader>,
       do_parse!(
           records: many0!(keyword_record) >>
               end_record >>
               many0!(blank_record) >>
               (PrimaryHeader::new(records))
       ));

named!(keyword_record<&[u8], KeywordRecord>,
       do_parse!(
           key: keyword  >>
               tag!("= ") >>
           vc: valuecomment >>
               (KeywordRecord::new(key, vc.0, vc.1))
       ));

named!(keyword<&[u8], Keyword>,
       map_res!(
           map_res!(
               take!(8),
               str::from_utf8),
           Keyword::from_str
       ));

named!(valuecomment<&[u8], (Value, Option<&str>)>,
       flat_map!(
           take!(70),
           pair!(
               value,
               opt!(comment)
           )));

named!(value<&[u8], Value>,
       map!(
           map_res!(
               is_not!("/"), // TODO Differentiate on the possible value types
               str::from_utf8
           ),
           Value::Raw
       )
);

named!(character_string<&[u8], Value>,
       map!(
           map_res!(
               ws!(delimited!(
                   tag!("'"),
                   take_while!(is_allowed_in_character_string),
                   tag!("'")
               )),
               str::from_utf8
           ),
           Value::CharacterString
       ));

fn is_allowed_in_character_string(chr: u8) -> bool {
    is_restricted_ascii(chr) && chr != 39
}

named!(comment<&[u8], &str>,
       map_res!(
           do_parse!(
               tag!("/") >>
                   comment: take_while!(is_comment_character) >>
                   (comment)
           ),
           str::from_utf8
       ));

fn is_comment_character(chr: u8) -> bool {
    is_restricted_ascii(chr)
}

fn is_restricted_ascii(chr: u8) -> bool {
    32u8 <= chr && chr <= 126u8
}

named!(end_record<&[u8], Keyword>,
       map!(
           flat_map!(
               take!(80),
               pair!(tag!("END"), many0!(tag!(" ")))
           ),
           |_| { Keyword::END }
       ));

named!(blank_record<&[u8], BlankRecord>,
       map!(
           count!(tag!(" "), 80),
           |_| { BlankRecord }
       ));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};
    use super::{fits, primary_header, keyword_record, keyword, valuecomment, character_string, end_record, blank_record};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, f) => {
                assert_eq!(f, Fits::new(long_cadence_header()));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, h) => assert_eq!(h, long_cadence_header()),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    fn long_cadence_header<'a>() -> PrimaryHeader<'a> {
        PrimaryHeader::new(vec!(
            KeywordRecord::convenient(Keyword::SIMPLE, "                   T ",
                               Option::Some(" conforms to FITS standards                     ")),
            KeywordRecord::convenient(Keyword::BITPIX, "                   8 ",
                               Option::Some(" array data type                                ")),
            KeywordRecord::convenient(Keyword::NAXIS, "                   0 ",
                               Option::Some(" number of array dimensions                     ")),
            KeywordRecord::convenient(Keyword::EXTEND, "                   T ",
                               Option::Some(" file contains extensions                       ")),
            KeywordRecord::convenient(Keyword::NEXTEND, "                   2 ",
                               Option::Some(" number of standard extensions                  ")),
            KeywordRecord::convenient(Keyword::EXTNAME, "'PRIMARY '           ",
                               Option::Some(" name of extension                              ")),
            KeywordRecord::convenient(Keyword::EXTVER, "                   1 ",
                               Option::Some(" extension version number (not format version)  ")),
            KeywordRecord::convenient(Keyword::ORIGIN, "'Unofficial data product' ",
                               Option::Some(" institution responsible for creating this ")),
            KeywordRecord::convenient(Keyword::DATE, "'2017-03-08'         ",
                               Option::Some(" file creation date.                            ")),
            KeywordRecord::convenient(Keyword::CREATOR, "'kadenza '           ",
                               Option::Some(" pipeline job and program u                     ")),
            KeywordRecord::convenient(Keyword::PROCVER, "'2.1.dev '           ",
                               Option::Some(" SW version                                     ")),
            KeywordRecord::convenient(Keyword::FILEVER, "'0.0     '           ",
                               Option::Some(" file format version                            ")),
            KeywordRecord::convenient(Keyword::TIMVERSN, "'' ",
                               Option::Some(" OGIP memo number for file format                                 ")),
            KeywordRecord::convenient(Keyword::TELESCOP, "'Kepler  '           ",
                               Option::Some(" telescope                                      ")),
            KeywordRecord::convenient(Keyword::INSTRUME, "'Kepler Photometer'  ",
                               Option::Some(" detector type                                  ")),
            KeywordRecord::convenient(Keyword::OBJECT, "'EPIC 200164267'     ",
                               Option::Some(" string version of target id                    ")),
            KeywordRecord::convenient(Keyword::KEPLERID, "           200164267 ",
                               Option::Some(" unique Kepler target identifier                ")),
            KeywordRecord::convenient(Keyword::CHANNEL, "                  68 ",
                               Option::Some(" CCD channel                                    ")),
            KeywordRecord::convenient(Keyword::MODULE, "                  19 ",
                               Option::Some(" CCD module                                     ")),
            KeywordRecord::convenient(Keyword::OUTPUT, "                   4 ",
                               Option::Some(" CCD output                                     ")),
            KeywordRecord::convenient(Keyword::CAMPAIGN, "'' ",
                               Option::Some(" Observing campaign number                                        ")),
            KeywordRecord::convenient(Keyword::DATA_REL, "'' ",
                               Option::Some(" data release version number                                      ")),
            KeywordRecord::convenient(Keyword::OBSMODE, "'long cadence'       ",
                               Option::Some(" observing mode                                 ")),
            KeywordRecord::convenient(Keyword::MISSION, "'K2      '           ",
                               Option::Some(" Mission name                                   ")),
            KeywordRecord::convenient(Keyword::TTABLEID, "'' ",
                               Option::Some(" target table id                                                  ")),
            KeywordRecord::convenient(Keyword::RADESYS, "'ICRS    '           ",
                               Option::Some(" reference frame of celestial coordinates       ")),
            KeywordRecord::convenient(Keyword::RA_OBJ, "'' ",
                               Option::Some(" [deg] right ascension                                            ")),
            KeywordRecord::convenient(Keyword::DEC_OBJ, "'' ",
                               Option::Some(" [deg] declination                                                ")),
            KeywordRecord::convenient(Keyword::EQUINOX, "              2000.0 ",
                               Option::Some(" equinox of celestial coordinate system         ")),
            KeywordRecord::convenient(Keyword::PMRA, " ",
                               Option::Some(" [arcsec/yr] RA proper motion                                       ")),
            KeywordRecord::convenient(Keyword::PMDEC, " ",
                               Option::Some(" [arcsec/yr] Dec proper motion                                      ")),
            KeywordRecord::convenient(Keyword::PMTOTAL, " ",
                               Option::Some(" [arcsec/yr] total proper motion                                    ")),
            KeywordRecord::convenient(Keyword::PARALLAX, " ",
                               Option::Some(" [arcsec] parallax                                                  ")),
            KeywordRecord::convenient(Keyword::GLON, " ",
                               Option::Some(" [deg] galactic longitude                                           ")),
            KeywordRecord::convenient(Keyword::GLAT, " ",
                               Option::Some(" [deg] galactic latitude                                            ")),
            KeywordRecord::convenient(Keyword::GMAG, " ",
                               Option::Some(" [mag] SDSS g band magnitude                                        ")),
            KeywordRecord::convenient(Keyword::RMAG, " ",
                               Option::Some(" [mag] SDSS r band magnitude                                        ")),
            KeywordRecord::convenient(Keyword::IMAG, " ",
                               Option::Some(" [mag] SDSS i band magnitude                                        ")),
            KeywordRecord::convenient(Keyword::ZMAG, " ",
                               Option::Some(" [mag] SDSS z band magnitude                                        ")),
            KeywordRecord::convenient(Keyword::JMAG, " ",
                               Option::Some(" [mag] J band magnitude from 2MASS                                  ")),
            KeywordRecord::convenient(Keyword::HMAG, " ",
                               Option::Some(" [mag] H band magnitude from 2MASS                                  ")),
            KeywordRecord::convenient(Keyword::KMAG, " ",
                               Option::Some(" [mag] K band magnitude from 2MASS                                  ")),
            KeywordRecord::convenient(Keyword::KEPMAG, " ",
                               Option::Some(" [mag] Kepler magnitude (Kp)                                        ")),
            KeywordRecord::convenient(Keyword::GRCOLOR, " ",
                               Option::Some(" [mag] (g-r) color, SDSS bands                                      ")),
            KeywordRecord::convenient(Keyword::JKCOLOR, " ",
                               Option::Some(" [mag] (J-K) color, 2MASS bands                                     ")),
            KeywordRecord::convenient(Keyword::GKCOLOR, " ",
                               Option::Some(" [mag] (g-K) color, SDSS g - 2MASS K                                ")),
            KeywordRecord::convenient(Keyword::TEFF, " ",
                               Option::Some(" [K] Effective temperature                                          ")),
            KeywordRecord::convenient(Keyword::LOGG, " ",
                               Option::Some(" [cm/s2] log10 surface gravity                                      ")),
            KeywordRecord::convenient(Keyword::FEH, " ",
                               Option::Some(" [log10([Fe/H])]  metallicity                                       ")),
            KeywordRecord::convenient(Keyword::EBMINUSV, " ",
                               Option::Some(" [mag] E(B-V) reddening                                             ")),
            KeywordRecord::convenient(Keyword::AV, " ",
                               Option::Some(" [mag] A_v extinction                                               ")),
            KeywordRecord::convenient(Keyword::RADIUS, " ",
                               Option::Some(" [solar radii] stellar radius                                       ")),
            KeywordRecord::convenient(Keyword::TMINDEX, " ",
                               Option::Some(" unique 2MASS catalog ID                                            ")),
            KeywordRecord::convenient(Keyword::CHECKSUM, "'7k7A7h637h697h69'   ",
                               Option::Some(" HDU checksum updated 2017-03-08T02:47:56       ")),
            KeywordRecord::convenient(Keyword::DATASUM, "'0       '           ",
                               Option::Some(" data unit checksum updated 2017-03-08T02:47:56 ")),
        ))
    }

    #[test]
    fn keyword_record_should_parse_a_keyword_record(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = keyword_record(data);

        match result {
            IResult::Done(_,k) => {
                assert_eq!(k, KeywordRecord::new(
                    Keyword::OBJECT,
                    Value::Raw("'EPIC 200164267'     "),
                    Option::Some(" string version of target id                    ")
                ))
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn valuecomment_should_parse_a_valuecomment(){
        let data = "'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = valuecomment(data);

        match result {
            IResult::Done(_, (value, comment)) => {
                assert_eq!(value, Value::Raw("'EPIC 200164267'     "));
                assert_eq!(comment, Option::Some(" string version of target id                    "));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn character_string_should_parse_an_quote_delimited_string(){
        let data = "   'EPIC 200164267'   "
            .as_bytes();

        let result = character_string(data);

        match result {
            IResult::Done(_, value) => {
                assert_eq!(value, Value::CharacterString("EPIC 200164267"));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_should_parse_a_keyword(){
        let data = "OBJECT  "
            .as_bytes();

        let result = keyword(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn end_record_should_parse_an_END_record(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_record(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::END),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn blank_record_should_parse_a_BLANK_record(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_record(data);

        match result {
            IResult::Done(_, record) => assert_eq!(record, BlankRecord),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<commit_msg>Introduce a convenient function to construct KeywordRecords<commit_after>//! The parser module is responsible for parsing FITS files.

use std::str;
use std::str::FromStr;
use super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};

named!(#[doc = "Will parse data from a FITS file into a `Fits` structure"], pub fits<&[u8], Fits>,
       do_parse!(
           ph: primary_header >>
               many0!(take!(2880)) >>
               (Fits::new(ph))
       ));

named!(primary_header<&[u8], PrimaryHeader>,
       do_parse!(
           records: many0!(keyword_record) >>
               end_record >>
               many0!(blank_record) >>
               (PrimaryHeader::new(records))
       ));

named!(keyword_record<&[u8], KeywordRecord>,
       do_parse!(
           key: keyword  >>
               tag!("= ") >>
           vc: valuecomment >>
               (KeywordRecord::new(key, vc.0, vc.1))
       ));

named!(keyword<&[u8], Keyword>,
       map_res!(
           map_res!(
               take!(8),
               str::from_utf8),
           Keyword::from_str
       ));

named!(valuecomment<&[u8], (Value, Option<&str>)>,
       flat_map!(
           take!(70),
           pair!(
               value,
               opt!(comment)
           )));

named!(value<&[u8], Value>,
       map!(
           map_res!(
               is_not!("/"), // TODO Differentiate on the possible value types
               str::from_utf8
           ),
           Value::Raw
       )
);

named!(character_string<&[u8], Value>,
       map!(
           map_res!(
               ws!(delimited!(
                   tag!("'"),
                   take_while!(is_allowed_in_character_string),
                   tag!("'")
               )),
               str::from_utf8
           ),
           Value::CharacterString
       ));

fn is_allowed_in_character_string(chr: u8) -> bool {
    is_restricted_ascii(chr) && chr != 39
}

named!(comment<&[u8], &str>,
       map_res!(
           do_parse!(
               tag!("/") >>
                   comment: take_while!(is_comment_character) >>
                   (comment)
           ),
           str::from_utf8
       ));

fn is_comment_character(chr: u8) -> bool {
    is_restricted_ascii(chr)
}

fn is_restricted_ascii(chr: u8) -> bool {
    32u8 <= chr && chr <= 126u8
}

named!(end_record<&[u8], Keyword>,
       map!(
           flat_map!(
               take!(80),
               pair!(tag!("END"), many0!(tag!(" ")))
           ),
           |_| { Keyword::END }
       ));

named!(blank_record<&[u8], BlankRecord>,
       map!(
           count!(tag!(" "), 80),
           |_| { BlankRecord }
       ));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};
    use super::{fits, primary_header, keyword_record, keyword, valuecomment, character_string, end_record, blank_record};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, f) => {
                assert_eq!(f, Fits::new(long_cadence_header()));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, h) => assert_eq!(h, long_cadence_header()),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    fn long_cadence_header<'a>() -> PrimaryHeader<'a> {
        PrimaryHeader::new(vec!(
            convenient(Keyword::SIMPLE, "                   T ",
                               Option::Some(" conforms to FITS standards                     ")),
            convenient(Keyword::BITPIX, "                   8 ",
                               Option::Some(" array data type                                ")),
            convenient(Keyword::NAXIS, "                   0 ",
                               Option::Some(" number of array dimensions                     ")),
            convenient(Keyword::EXTEND, "                   T ",
                               Option::Some(" file contains extensions                       ")),
            convenient(Keyword::NEXTEND, "                   2 ",
                               Option::Some(" number of standard extensions                  ")),
            convenient(Keyword::EXTNAME, "'PRIMARY '           ",
                               Option::Some(" name of extension                              ")),
            convenient(Keyword::EXTVER, "                   1 ",
                               Option::Some(" extension version number (not format version)  ")),
            convenient(Keyword::ORIGIN, "'Unofficial data product' ",
                               Option::Some(" institution responsible for creating this ")),
            convenient(Keyword::DATE, "'2017-03-08'         ",
                               Option::Some(" file creation date.                            ")),
            convenient(Keyword::CREATOR, "'kadenza '           ",
                               Option::Some(" pipeline job and program u                     ")),
            convenient(Keyword::PROCVER, "'2.1.dev '           ",
                               Option::Some(" SW version                                     ")),
            convenient(Keyword::FILEVER, "'0.0     '           ",
                               Option::Some(" file format version                            ")),
            convenient(Keyword::TIMVERSN, "'' ",
                               Option::Some(" OGIP memo number for file format                                 ")),
            convenient(Keyword::TELESCOP, "'Kepler  '           ",
                               Option::Some(" telescope                                      ")),
            convenient(Keyword::INSTRUME, "'Kepler Photometer'  ",
                               Option::Some(" detector type                                  ")),
            convenient(Keyword::OBJECT, "'EPIC 200164267'     ",
                               Option::Some(" string version of target id                    ")),
            convenient(Keyword::KEPLERID, "           200164267 ",
                               Option::Some(" unique Kepler target identifier                ")),
            convenient(Keyword::CHANNEL, "                  68 ",
                               Option::Some(" CCD channel                                    ")),
            convenient(Keyword::MODULE, "                  19 ",
                               Option::Some(" CCD module                                     ")),
            convenient(Keyword::OUTPUT, "                   4 ",
                               Option::Some(" CCD output                                     ")),
            convenient(Keyword::CAMPAIGN, "'' ",
                               Option::Some(" Observing campaign number                                        ")),
            convenient(Keyword::DATA_REL, "'' ",
                               Option::Some(" data release version number                                      ")),
            convenient(Keyword::OBSMODE, "'long cadence'       ",
                               Option::Some(" observing mode                                 ")),
            convenient(Keyword::MISSION, "'K2      '           ",
                               Option::Some(" Mission name                                   ")),
            convenient(Keyword::TTABLEID, "'' ",
                               Option::Some(" target table id                                                  ")),
            convenient(Keyword::RADESYS, "'ICRS    '           ",
                               Option::Some(" reference frame of celestial coordinates       ")),
            convenient(Keyword::RA_OBJ, "'' ",
                               Option::Some(" [deg] right ascension                                            ")),
            convenient(Keyword::DEC_OBJ, "'' ",
                               Option::Some(" [deg] declination                                                ")),
            convenient(Keyword::EQUINOX, "              2000.0 ",
                               Option::Some(" equinox of celestial coordinate system         ")),
            convenient(Keyword::PMRA, " ",
                               Option::Some(" [arcsec/yr] RA proper motion                                       ")),
            convenient(Keyword::PMDEC, " ",
                               Option::Some(" [arcsec/yr] Dec proper motion                                      ")),
            convenient(Keyword::PMTOTAL, " ",
                               Option::Some(" [arcsec/yr] total proper motion                                    ")),
            convenient(Keyword::PARALLAX, " ",
                               Option::Some(" [arcsec] parallax                                                  ")),
            convenient(Keyword::GLON, " ",
                               Option::Some(" [deg] galactic longitude                                           ")),
            convenient(Keyword::GLAT, " ",
                               Option::Some(" [deg] galactic latitude                                            ")),
            convenient(Keyword::GMAG, " ",
                               Option::Some(" [mag] SDSS g band magnitude                                        ")),
            convenient(Keyword::RMAG, " ",
                               Option::Some(" [mag] SDSS r band magnitude                                        ")),
            convenient(Keyword::IMAG, " ",
                               Option::Some(" [mag] SDSS i band magnitude                                        ")),
            convenient(Keyword::ZMAG, " ",
                               Option::Some(" [mag] SDSS z band magnitude                                        ")),
            convenient(Keyword::JMAG, " ",
                               Option::Some(" [mag] J band magnitude from 2MASS                                  ")),
            convenient(Keyword::HMAG, " ",
                               Option::Some(" [mag] H band magnitude from 2MASS                                  ")),
            convenient(Keyword::KMAG, " ",
                               Option::Some(" [mag] K band magnitude from 2MASS                                  ")),
            convenient(Keyword::KEPMAG, " ",
                               Option::Some(" [mag] Kepler magnitude (Kp)                                        ")),
            convenient(Keyword::GRCOLOR, " ",
                               Option::Some(" [mag] (g-r) color, SDSS bands                                      ")),
            convenient(Keyword::JKCOLOR, " ",
                               Option::Some(" [mag] (J-K) color, 2MASS bands                                     ")),
            convenient(Keyword::GKCOLOR, " ",
                               Option::Some(" [mag] (g-K) color, SDSS g - 2MASS K                                ")),
            convenient(Keyword::TEFF, " ",
                               Option::Some(" [K] Effective temperature                                          ")),
            convenient(Keyword::LOGG, " ",
                               Option::Some(" [cm/s2] log10 surface gravity                                      ")),
            convenient(Keyword::FEH, " ",
                               Option::Some(" [log10([Fe/H])]  metallicity                                       ")),
            convenient(Keyword::EBMINUSV, " ",
                               Option::Some(" [mag] E(B-V) reddening                                             ")),
            convenient(Keyword::AV, " ",
                               Option::Some(" [mag] A_v extinction                                               ")),
            convenient(Keyword::RADIUS, " ",
                               Option::Some(" [solar radii] stellar radius                                       ")),
            convenient(Keyword::TMINDEX, " ",
                               Option::Some(" unique 2MASS catalog ID                                            ")),
            convenient(Keyword::CHECKSUM, "'7k7A7h637h697h69'   ",
                               Option::Some(" HDU checksum updated 2017-03-08T02:47:56       ")),
            convenient(Keyword::DATASUM, "'0       '           ",
                               Option::Some(" data unit checksum updated 2017-03-08T02:47:56 ")),
        ))
    }

    fn convenient<'a>(keyword: Keyword, value: &'a str, comment: Option<&'a str>) -> KeywordRecord<'a> {
        KeywordRecord::convenient(keyword, value, comment)
    }

    #[test]
    fn keyword_record_should_parse_a_keyword_record(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = keyword_record(data);

        match result {
            IResult::Done(_,k) => {
                assert_eq!(k, KeywordRecord::new(
                    Keyword::OBJECT,
                    Value::Raw("'EPIC 200164267'     "),
                    Option::Some(" string version of target id                    ")
                ))
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn valuecomment_should_parse_a_valuecomment(){
        let data = "'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = valuecomment(data);

        match result {
            IResult::Done(_, (value, comment)) => {
                assert_eq!(value, Value::Raw("'EPIC 200164267'     "));
                assert_eq!(comment, Option::Some(" string version of target id                    "));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn character_string_should_parse_an_quote_delimited_string(){
        let data = "   'EPIC 200164267'   "
            .as_bytes();

        let result = character_string(data);

        match result {
            IResult::Done(_, value) => {
                assert_eq!(value, Value::CharacterString("EPIC 200164267"));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_should_parse_a_keyword(){
        let data = "OBJECT  "
            .as_bytes();

        let result = keyword(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn end_record_should_parse_an_END_record(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_record(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::END),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn blank_record_should_parse_a_BLANK_record(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_record(data);

        match result {
            IResult::Done(_, record) => assert_eq!(record, BlankRecord),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<|endoftext|>


<commit_before>//! The parser module is responsible for parsing FITS files.

use std::str;
use std::str::FromStr;
use super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};

named!(#[doc = "Will parse data from a FITS file into a `Fits` structure"], pub fits<&[u8], Fits>,
       do_parse!(
           ph: primary_header >>
               many0!(take!(2880)) >>
               (Fits::new(ph))
       ));

named!(primary_header<&[u8], PrimaryHeader>,
       do_parse!(
           records: many0!(keyword_record) >>
               end_record >>
               many0!(blank_record) >>
               (PrimaryHeader::new(records))
       ));

named!(keyword_record<&[u8], KeywordRecord>,
       do_parse!(
           key: keyword  >>
               tag!("= ") >>
           vc: valuecomment >>
               (KeywordRecord::new(key, vc.0, vc.1))
       ));

named!(keyword<&[u8], Keyword>,
       map_res!(
           map_res!(
               take!(8),
               str::from_utf8),
           Keyword::from_str
       ));

named!(valuecomment<&[u8], (Value, Option<&str>)>,
       flat_map!(
           take!(70),
           pair!(
               value,
               opt!(comment)
           )));

named!(value<&[u8], Value>,
       map!(
           map_res!(
               is_not!("/"), // TODO Differentiate on the possible value types
               str::from_utf8
           ),
           Value::Raw
       )
);

named!(character_string<&[u8], Value>,
       map!(
           map_res!(
               ws!(delimited!(
                   tag!("'"),
                   take_while!(is_allowed_in_character_string),
                   tag!("'")
               )),
               str::from_utf8
           ),
           Value::CharacterString
       ));

fn is_allowed_in_character_string(chr: u8) -> bool {
    is_restricted_ascii(chr) && chr != 39
}

named!(comment<&[u8], &str>,
       map_res!(
           do_parse!(
               tag!("/") >>
                   comment: take_while!(is_comment_character) >>
                   (comment)
           ),
           str::from_utf8
       ));

fn is_comment_character(chr: u8) -> bool {
    is_restricted_ascii(chr)
}

fn is_restricted_ascii(chr: u8) -> bool {
    32u8 <= chr && chr <= 126u8
}

named!(end_record<&[u8], Keyword>,
       map!(
           flat_map!(
               take!(80),
               pair!(tag!("END"), many0!(tag!(" ")))
           ),
           |_| { Keyword::END }
       ));

named!(blank_record<&[u8], BlankRecord>,
       map!(
           count!(tag!(" "), 80),
           |_| { BlankRecord }
       ));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};
    use super::{fits, primary_header, keyword_record, keyword, valuecomment, character_string, end_record, blank_record};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, f) => {
                assert_eq!(f, Fits::new(long_cadence_header()));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, h) => assert_eq!(h, long_cadence_header()),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    fn long_cadence_header<'a>() -> PrimaryHeader<'a> {
        PrimaryHeader::new(vec!(
            convenient(Keyword::SIMPLE, "                   T ",
                               Option::Some(" conforms to FITS standards                     ")),
            convenient(Keyword::BITPIX, "                   8 ",
                               Option::Some(" array data type                                ")),
            convenient(Keyword::NAXIS, "                   0 ",
                               Option::Some(" number of array dimensions                     ")),
            convenient(Keyword::EXTEND, "                   T ",
                               Option::Some(" file contains extensions                       ")),
            convenient(Keyword::NEXTEND, "                   2 ",
                               Option::Some(" number of standard extensions                  ")),
            convenient(Keyword::EXTNAME, "'PRIMARY '           ",
                               Option::Some(" name of extension                              ")),
            convenient(Keyword::EXTVER, "                   1 ",
                               Option::Some(" extension version number (not format version)  ")),
            convenient(Keyword::ORIGIN, "'Unofficial data product' ",
                               Option::Some(" institution responsible for creating this ")),
            convenient(Keyword::DATE, "'2017-03-08'         ",
                               Option::Some(" file creation date.                            ")),
            convenient(Keyword::CREATOR, "'kadenza '           ",
                               Option::Some(" pipeline job and program u                     ")),
            convenient(Keyword::PROCVER, "'2.1.dev '           ",
                               Option::Some(" SW version                                     ")),
            convenient(Keyword::FILEVER, "'0.0     '           ",
                               Option::Some(" file format version                            ")),
            convenient(Keyword::TIMVERSN, "'' ",
                               Option::Some(" OGIP memo number for file format                                 ")),
            convenient(Keyword::TELESCOP, "'Kepler  '           ",
                               Option::Some(" telescope                                      ")),
            convenient(Keyword::INSTRUME, "'Kepler Photometer'  ",
                               Option::Some(" detector type                                  ")),
            convenient(Keyword::OBJECT, "'EPIC 200164267'     ",
                               Option::Some(" string version of target id                    ")),
            convenient(Keyword::KEPLERID, "           200164267 ",
                               Option::Some(" unique Kepler target identifier                ")),
            convenient(Keyword::CHANNEL, "                  68 ",
                               Option::Some(" CCD channel                                    ")),
            convenient(Keyword::MODULE, "                  19 ",
                               Option::Some(" CCD module                                     ")),
            convenient(Keyword::OUTPUT, "                   4 ",
                               Option::Some(" CCD output                                     ")),
            convenient(Keyword::CAMPAIGN, "'' ",
                               Option::Some(" Observing campaign number                                        ")),
            convenient(Keyword::DATA_REL, "'' ",
                               Option::Some(" data release version number                                      ")),
            convenient(Keyword::OBSMODE, "'long cadence'       ",
                               Option::Some(" observing mode                                 ")),
            convenient(Keyword::MISSION, "'K2      '           ",
                               Option::Some(" Mission name                                   ")),
            convenient(Keyword::TTABLEID, "'' ",
                               Option::Some(" target table id                                                  ")),
            convenient(Keyword::RADESYS, "'ICRS    '           ",
                               Option::Some(" reference frame of celestial coordinates       ")),
            convenient(Keyword::RA_OBJ, "'' ",
                               Option::Some(" [deg] right ascension                                            ")),
            convenient(Keyword::DEC_OBJ, "'' ",
                               Option::Some(" [deg] declination                                                ")),
            convenient(Keyword::EQUINOX, "              2000.0 ",
                               Option::Some(" equinox of celestial coordinate system         ")),
            convenient(Keyword::PMRA, " ",
                               Option::Some(" [arcsec/yr] RA proper motion                                       ")),
            convenient(Keyword::PMDEC, " ",
                               Option::Some(" [arcsec/yr] Dec proper motion                                      ")),
            convenient(Keyword::PMTOTAL, " ",
                               Option::Some(" [arcsec/yr] total proper motion                                    ")),
            convenient(Keyword::PARALLAX, " ",
                               Option::Some(" [arcsec] parallax                                                  ")),
            convenient(Keyword::GLON, " ",
                               Option::Some(" [deg] galactic longitude                                           ")),
            convenient(Keyword::GLAT, " ",
                               Option::Some(" [deg] galactic latitude                                            ")),
            convenient(Keyword::GMAG, " ",
                               Option::Some(" [mag] SDSS g band magnitude                                        ")),
            convenient(Keyword::RMAG, " ",
                               Option::Some(" [mag] SDSS r band magnitude                                        ")),
            convenient(Keyword::IMAG, " ",
                               Option::Some(" [mag] SDSS i band magnitude                                        ")),
            convenient(Keyword::ZMAG, " ",
                               Option::Some(" [mag] SDSS z band magnitude                                        ")),
            convenient(Keyword::JMAG, " ",
                               Option::Some(" [mag] J band magnitude from 2MASS                                  ")),
            convenient(Keyword::HMAG, " ",
                               Option::Some(" [mag] H band magnitude from 2MASS                                  ")),
            convenient(Keyword::KMAG, " ",
                               Option::Some(" [mag] K band magnitude from 2MASS                                  ")),
            convenient(Keyword::KEPMAG, " ",
                               Option::Some(" [mag] Kepler magnitude (Kp)                                        ")),
            convenient(Keyword::GRCOLOR, " ",
                               Option::Some(" [mag] (g-r) color, SDSS bands                                      ")),
            convenient(Keyword::JKCOLOR, " ",
                               Option::Some(" [mag] (J-K) color, 2MASS bands                                     ")),
            convenient(Keyword::GKCOLOR, " ",
                               Option::Some(" [mag] (g-K) color, SDSS g - 2MASS K                                ")),
            convenient(Keyword::TEFF, " ",
                               Option::Some(" [K] Effective temperature                                          ")),
            convenient(Keyword::LOGG, " ",
                               Option::Some(" [cm/s2] log10 surface gravity                                      ")),
            convenient(Keyword::FEH, " ",
                               Option::Some(" [log10([Fe/H])]  metallicity                                       ")),
            convenient(Keyword::EBMINUSV, " ",
                               Option::Some(" [mag] E(B-V) reddening                                             ")),
            convenient(Keyword::AV, " ",
                               Option::Some(" [mag] A_v extinction                                               ")),
            convenient(Keyword::RADIUS, " ",
                               Option::Some(" [solar radii] stellar radius                                       ")),
            convenient(Keyword::TMINDEX, " ",
                               Option::Some(" unique 2MASS catalog ID                                            ")),
            convenient(Keyword::CHECKSUM, "'7k7A7h637h697h69'   ",
                               Option::Some(" HDU checksum updated 2017-03-08T02:47:56       ")),
            convenient(Keyword::DATASUM, "'0       '           ",
                               Option::Some(" data unit checksum updated 2017-03-08T02:47:56 ")),
        ))
    }

    fn convenient<'a>(keyword: Keyword, value: &'a str, comment: Option<&'a str>) -> KeywordRecord<'a> {
        KeywordRecord::new(keyword, Value::Raw(value), comment)
    }

    #[test]
    fn keyword_record_should_parse_a_keyword_record(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = keyword_record(data);

        match result {
            IResult::Done(_,k) => {
                assert_eq!(k, KeywordRecord::new(
                    Keyword::OBJECT,
                    Value::Raw("'EPIC 200164267'     "),
                    Option::Some(" string version of target id                    ")
                ))
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn valuecomment_should_parse_a_valuecomment(){
        let data = "'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = valuecomment(data);

        match result {
            IResult::Done(_, (value, comment)) => {
                assert_eq!(value, Value::Raw("'EPIC 200164267'     "));
                assert_eq!(comment, Option::Some(" string version of target id                    "));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn character_string_should_parse_an_quote_delimited_string(){
        let data = "   'EPIC 200164267'   "
            .as_bytes();

        let result = character_string(data);

        match result {
            IResult::Done(_, value) => {
                assert_eq!(value, Value::CharacterString("EPIC 200164267"));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_should_parse_a_keyword(){
        let data = "OBJECT  "
            .as_bytes();

        let result = keyword(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn end_record_should_parse_an_END_record(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_record(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::END),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn blank_record_should_parse_a_BLANK_record(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_record(data);

        match result {
            IResult::Done(_, record) => assert_eq!(record, BlankRecord),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<commit_msg>Remove is_comment_character in favor of is_restricted_ascii<commit_after>//! The parser module is responsible for parsing FITS files.

use std::str;
use std::str::FromStr;
use super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};

named!(#[doc = "Will parse data from a FITS file into a `Fits` structure"], pub fits<&[u8], Fits>,
       do_parse!(
           ph: primary_header >>
               many0!(take!(2880)) >>
               (Fits::new(ph))
       ));

named!(primary_header<&[u8], PrimaryHeader>,
       do_parse!(
           records: many0!(keyword_record) >>
               end_record >>
               many0!(blank_record) >>
               (PrimaryHeader::new(records))
       ));

named!(keyword_record<&[u8], KeywordRecord>,
       do_parse!(
           key: keyword  >>
               tag!("= ") >>
           vc: valuecomment >>
               (KeywordRecord::new(key, vc.0, vc.1))
       ));

named!(keyword<&[u8], Keyword>,
       map_res!(
           map_res!(
               take!(8),
               str::from_utf8),
           Keyword::from_str
       ));

named!(valuecomment<&[u8], (Value, Option<&str>)>,
       flat_map!(
           take!(70),
           pair!(
               value,
               opt!(comment)
           )));

named!(value<&[u8], Value>,
       map!(
           map_res!(
               is_not!("/"), // TODO Differentiate on the possible value types
               str::from_utf8
           ),
           Value::Raw
       )
);

named!(character_string<&[u8], Value>,
       map!(
           map_res!(
               ws!(delimited!(
                   tag!("'"),
                   take_while!(is_allowed_in_character_string),
                   tag!("'")
               )),
               str::from_utf8
           ),
           Value::CharacterString
       ));

fn is_allowed_in_character_string(chr: u8) -> bool {
    is_restricted_ascii(chr) && chr != 39
}

named!(comment<&[u8], &str>,
       map_res!(
           do_parse!(
               tag!("/") >>
                   comment: take_while!(is_restricted_ascii) >>
                   (comment)
           ),
           str::from_utf8
       ));

fn is_restricted_ascii(chr: u8) -> bool {
    32u8 <= chr && chr <= 126u8
}

named!(end_record<&[u8], Keyword>,
       map!(
           flat_map!(
               take!(80),
               pair!(tag!("END"), many0!(tag!(" ")))
           ),
           |_| { Keyword::END }
       ));

named!(blank_record<&[u8], BlankRecord>,
       map!(
           count!(tag!(" "), 80),
           |_| { BlankRecord }
       ));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::super::types::{Fits, PrimaryHeader, KeywordRecord, Keyword, Value, BlankRecord};
    use super::{fits, primary_header, keyword_record, keyword, valuecomment, character_string, end_record, blank_record};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, f) => {
                assert_eq!(f, Fits::new(long_cadence_header()));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, h) => assert_eq!(h, long_cadence_header()),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    fn long_cadence_header<'a>() -> PrimaryHeader<'a> {
        PrimaryHeader::new(vec!(
            convenient(Keyword::SIMPLE, "                   T ",
                               Option::Some(" conforms to FITS standards                     ")),
            convenient(Keyword::BITPIX, "                   8 ",
                               Option::Some(" array data type                                ")),
            convenient(Keyword::NAXIS, "                   0 ",
                               Option::Some(" number of array dimensions                     ")),
            convenient(Keyword::EXTEND, "                   T ",
                               Option::Some(" file contains extensions                       ")),
            convenient(Keyword::NEXTEND, "                   2 ",
                               Option::Some(" number of standard extensions                  ")),
            convenient(Keyword::EXTNAME, "'PRIMARY '           ",
                               Option::Some(" name of extension                              ")),
            convenient(Keyword::EXTVER, "                   1 ",
                               Option::Some(" extension version number (not format version)  ")),
            convenient(Keyword::ORIGIN, "'Unofficial data product' ",
                               Option::Some(" institution responsible for creating this ")),
            convenient(Keyword::DATE, "'2017-03-08'         ",
                               Option::Some(" file creation date.                            ")),
            convenient(Keyword::CREATOR, "'kadenza '           ",
                               Option::Some(" pipeline job and program u                     ")),
            convenient(Keyword::PROCVER, "'2.1.dev '           ",
                               Option::Some(" SW version                                     ")),
            convenient(Keyword::FILEVER, "'0.0     '           ",
                               Option::Some(" file format version                            ")),
            convenient(Keyword::TIMVERSN, "'' ",
                               Option::Some(" OGIP memo number for file format                                 ")),
            convenient(Keyword::TELESCOP, "'Kepler  '           ",
                               Option::Some(" telescope                                      ")),
            convenient(Keyword::INSTRUME, "'Kepler Photometer'  ",
                               Option::Some(" detector type                                  ")),
            convenient(Keyword::OBJECT, "'EPIC 200164267'     ",
                               Option::Some(" string version of target id                    ")),
            convenient(Keyword::KEPLERID, "           200164267 ",
                               Option::Some(" unique Kepler target identifier                ")),
            convenient(Keyword::CHANNEL, "                  68 ",
                               Option::Some(" CCD channel                                    ")),
            convenient(Keyword::MODULE, "                  19 ",
                               Option::Some(" CCD module                                     ")),
            convenient(Keyword::OUTPUT, "                   4 ",
                               Option::Some(" CCD output                                     ")),
            convenient(Keyword::CAMPAIGN, "'' ",
                               Option::Some(" Observing campaign number                                        ")),
            convenient(Keyword::DATA_REL, "'' ",
                               Option::Some(" data release version number                                      ")),
            convenient(Keyword::OBSMODE, "'long cadence'       ",
                               Option::Some(" observing mode                                 ")),
            convenient(Keyword::MISSION, "'K2      '           ",
                               Option::Some(" Mission name                                   ")),
            convenient(Keyword::TTABLEID, "'' ",
                               Option::Some(" target table id                                                  ")),
            convenient(Keyword::RADESYS, "'ICRS    '           ",
                               Option::Some(" reference frame of celestial coordinates       ")),
            convenient(Keyword::RA_OBJ, "'' ",
                               Option::Some(" [deg] right ascension                                            ")),
            convenient(Keyword::DEC_OBJ, "'' ",
                               Option::Some(" [deg] declination                                                ")),
            convenient(Keyword::EQUINOX, "              2000.0 ",
                               Option::Some(" equinox of celestial coordinate system         ")),
            convenient(Keyword::PMRA, " ",
                               Option::Some(" [arcsec/yr] RA proper motion                                       ")),
            convenient(Keyword::PMDEC, " ",
                               Option::Some(" [arcsec/yr] Dec proper motion                                      ")),
            convenient(Keyword::PMTOTAL, " ",
                               Option::Some(" [arcsec/yr] total proper motion                                    ")),
            convenient(Keyword::PARALLAX, " ",
                               Option::Some(" [arcsec] parallax                                                  ")),
            convenient(Keyword::GLON, " ",
                               Option::Some(" [deg] galactic longitude                                           ")),
            convenient(Keyword::GLAT, " ",
                               Option::Some(" [deg] galactic latitude                                            ")),
            convenient(Keyword::GMAG, " ",
                               Option::Some(" [mag] SDSS g band magnitude                                        ")),
            convenient(Keyword::RMAG, " ",
                               Option::Some(" [mag] SDSS r band magnitude                                        ")),
            convenient(Keyword::IMAG, " ",
                               Option::Some(" [mag] SDSS i band magnitude                                        ")),
            convenient(Keyword::ZMAG, " ",
                               Option::Some(" [mag] SDSS z band magnitude                                        ")),
            convenient(Keyword::JMAG, " ",
                               Option::Some(" [mag] J band magnitude from 2MASS                                  ")),
            convenient(Keyword::HMAG, " ",
                               Option::Some(" [mag] H band magnitude from 2MASS                                  ")),
            convenient(Keyword::KMAG, " ",
                               Option::Some(" [mag] K band magnitude from 2MASS                                  ")),
            convenient(Keyword::KEPMAG, " ",
                               Option::Some(" [mag] Kepler magnitude (Kp)                                        ")),
            convenient(Keyword::GRCOLOR, " ",
                               Option::Some(" [mag] (g-r) color, SDSS bands                                      ")),
            convenient(Keyword::JKCOLOR, " ",
                               Option::Some(" [mag] (J-K) color, 2MASS bands                                     ")),
            convenient(Keyword::GKCOLOR, " ",
                               Option::Some(" [mag] (g-K) color, SDSS g - 2MASS K                                ")),
            convenient(Keyword::TEFF, " ",
                               Option::Some(" [K] Effective temperature                                          ")),
            convenient(Keyword::LOGG, " ",
                               Option::Some(" [cm/s2] log10 surface gravity                                      ")),
            convenient(Keyword::FEH, " ",
                               Option::Some(" [log10([Fe/H])]  metallicity                                       ")),
            convenient(Keyword::EBMINUSV, " ",
                               Option::Some(" [mag] E(B-V) reddening                                             ")),
            convenient(Keyword::AV, " ",
                               Option::Some(" [mag] A_v extinction                                               ")),
            convenient(Keyword::RADIUS, " ",
                               Option::Some(" [solar radii] stellar radius                                       ")),
            convenient(Keyword::TMINDEX, " ",
                               Option::Some(" unique 2MASS catalog ID                                            ")),
            convenient(Keyword::CHECKSUM, "'7k7A7h637h697h69'   ",
                               Option::Some(" HDU checksum updated 2017-03-08T02:47:56       ")),
            convenient(Keyword::DATASUM, "'0       '           ",
                               Option::Some(" data unit checksum updated 2017-03-08T02:47:56 ")),
        ))
    }

    fn convenient<'a>(keyword: Keyword, value: &'a str, comment: Option<&'a str>) -> KeywordRecord<'a> {
        KeywordRecord::new(keyword, Value::Raw(value), comment)
    }

    #[test]
    fn keyword_record_should_parse_a_keyword_record(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = keyword_record(data);

        match result {
            IResult::Done(_,k) => {
                assert_eq!(k, KeywordRecord::new(
                    Keyword::OBJECT,
                    Value::Raw("'EPIC 200164267'     "),
                    Option::Some(" string version of target id                    ")
                ))
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn valuecomment_should_parse_a_valuecomment(){
        let data = "'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = valuecomment(data);

        match result {
            IResult::Done(_, (value, comment)) => {
                assert_eq!(value, Value::Raw("'EPIC 200164267'     "));
                assert_eq!(comment, Option::Some(" string version of target id                    "));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn character_string_should_parse_an_quote_delimited_string(){
        let data = "   'EPIC 200164267'   "
            .as_bytes();

        let result = character_string(data);

        match result {
            IResult::Done(_, value) => {
                assert_eq!(value, Value::CharacterString("EPIC 200164267"));
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_should_parse_a_keyword(){
        let data = "OBJECT  "
            .as_bytes();

        let result = keyword(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn end_record_should_parse_an_END_record(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_record(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::END),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    #[allow(non_snake_case)]
    fn blank_record_should_parse_a_BLANK_record(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_record(data);

        match result {
            IResult::Done(_, record) => assert_eq!(record, BlankRecord),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<|endoftext|>


<commit_before>use super::fitness::Fitness;
use super::traits::{Genotype, Distance};
use rand::Rng;
use closed01::Closed01;
use super::traits::Mate;
use std::marker::PhantomData;
use std::num::Zero;
use std::fmt::Debug;
use std::cmp;
use std::mem;
use rayon::par_iter::*;
use super::prob::probabilistic_round;

#[derive(Debug)]
pub struct Individual<T: Debug + Genotype> {
    fitness: Fitness,
    genome: Box<T>,
}

impl<T: Debug + Genotype> Individual<T> {
    pub fn fitness(&self) -> Fitness {
        self.fitness
    }

    pub fn genome(&self) -> &T {
        &self.genome
    }
}

pub trait Rating { }

pub trait IsRated : Rating { }
pub trait IsRatedSorted : IsRated { }

#[derive(Debug)]
pub struct Unrated;

#[derive(Debug)]
pub struct Rated;

#[derive(Debug)]
pub struct RatedSorted;

impl Rating for Unrated {}
impl Rating for Rated {}
impl Rating for RatedSorted {}

impl IsRated for Rated {}
impl IsRated for RatedSorted {}
impl IsRatedSorted for RatedSorted {}

#[derive(Debug)]
pub struct Population<T: Genotype + Debug, R: Rating> {
    individuals: Vec<Individual<T>>,
    _marker: PhantomData<R>,
}

#[derive(Debug)]
pub struct Niches<T: Genotype + Debug> {
    total_individuals: usize,
    niches: Vec<Population<T, Rated>>,
}

impl<T: Genotype + Debug> Niches<T> {
    pub fn new() -> Self {
        Niches {
            total_individuals: 0,
            niches: Vec::new(),
        }
    }

    pub fn into_iter(self) -> ::std::vec::IntoIter<Population<T, Rated>> {
        self.niches.into_iter()
    }

    /// The sum of all "mean fitnesses" of all niches.

    fn total_mean(&self) -> Fitness {
        self.niches.iter().map(|ind| ind.mean_fitness()).sum()
    }

    /// Number of niches

    pub fn len(&self) -> usize {
        self.niches.len()
    }

    /// Add an individual to one of the niches.

    pub fn add_individual<R, C>(&mut self,
                                ind: Individual<T>,
                                compatibility_threshold: f64,
                                compatibility: &C,
                                rng: &mut R)
        where R: Rng,
              C: Distance<T>
    {
        self.total_individuals += 1;

        for niche in self.niches.iter_mut() {
            // Is this genome compatible with this niche? Test against a random genome.
            let compatible = match rng.choose(&niche.individuals) {
                Some(probe) => {
                    compatibility.distance(&probe.genome, &ind.genome) < compatibility_threshold
                }
                // If a niche is empty, a genome always is compatible (note that a niche can't be empyt)
                None => true,
            };
            if compatible {
                niche.add_individual(ind);
                return;
            }
        }
        // if no compatible niche was found, create a new niche containing this genome.
        let mut new_niche = Population::new();
        new_niche.add_individual(ind);
        self.niches.push(new_niche);
    }
}

impl<T: Genotype + Debug, R: Rating> Population<T, R> {
    pub fn new() -> Population<T, R> {
        Population {
            individuals: Vec::new(),
            _marker: PhantomData,
        }
    }

    pub fn new_from_vec(v: Vec<Individual<T>>) -> Population<T, R> {
        Population {
            individuals: v,
            _marker: PhantomData,
        }
    }

    pub fn len(&self) -> usize {
        self.individuals.len()
    }
}

impl<T: Genotype + Debug> Population<T, Unrated> {
    pub fn add_genome(&mut self, genome: Box<T>) {
        self.individuals.push(Individual {
            fitness: Zero::zero(),
            genome: genome,
        });
    }

    pub fn rate<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Fn(&T) -> Fitness
    {
        for ind in self.individuals.iter_mut() {
            let fitness = f(&ind.genome);
            ind.fitness = fitness;
        }
        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }

    pub fn rate_par<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Sync + Fn(&T) -> Fitness
    {
        self.individuals.par_iter_mut().for_each(|ind| {
            let fitness = f(&ind.genome);
            ind.fitness = fitness;
        });

        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }
}

impl<T: Genotype + Debug, R: IsRated> Population<T, R> {
    fn mean_fitness(&self) -> Fitness {
        let sum: Fitness = self.individuals.iter().map(|ind| ind.fitness).sum();
        sum / Fitness::new(self.len() as f64)
    }
}

impl<T: Genotype + Debug> Population<T, RatedSorted> {
    pub fn best_individual(&self) -> Option<&Individual<T>> {
        self.individuals.first()
    }

    // Return true if genome at position `i` is fitter that `j`
    //
    // In a sorted population, the individual with the lower index
    // has a better fitness.

    // #[inline]
    // fn is_fitter(&self, i: usize, j: usize) -> bool {
    //    i < j
    // }

    /// Create a single offspring Genome by selecting random parents
    /// from the best `select_size` individuals of the populations.

    fn create_single_offspring<R, M>(&self, select_size: usize, mate: &mut M, rng: &mut R) -> T
        where R: Rng,
              M: Mate<T>
    {
        assert!(select_size > 0 && select_size <= self.len());

        // We do not need tournament selection here as our population is sorted.
        // We simply determine two individuals out of `select_size`.

        let mut parent1 = rng.gen_range(0, select_size);
        let mut parent2 = rng.gen_range(0, select_size);

        // try to find a parent2 != parent1. retry three times.
        for _ in 0..3 {
            if parent2 != parent1 {
                break;
            }
            parent2 = rng.gen_range(0, select_size);
        }

        // `mate` assumes that the first parent performs better.
        if parent1 > parent2 {
            mem::swap(&mut parent1, &mut parent2);
        }

        debug_assert!(parent1 <= parent2);

        mate.mate(&self.individuals[parent1].genome,
                  &self.individuals[parent2].genome,
                  parent1 == parent2,
                  rng)
    }
}

impl<T: Genotype + Debug> Population<T, Rated> {
    pub fn add_individual(&mut self, ind: Individual<T>) {
        self.individuals.push(ind);
    }

    // Return true if genome at position `i` is fitter that `j`
    // fn is_fitter(&self, i: usize, j: usize) -> bool {
    //    self.individuals[i].fitness > self.individuals[j].fitness
    // }

    // higher value of fitness means that the individual is fitter.
    pub fn sort(mut self) -> Population<T, RatedSorted> {
        (&mut self.individuals).sort_by(|a, b| a.fitness.cmp(&b.fitness).reverse());
        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }

    pub fn best_individual(&self) -> Option<&Individual<T>> {
        self.individuals.iter().max_by_key(|ind| ind.fitness)
    }

    /// Merge `self` with the first `n` individuals from population `other`.
    pub fn merge(&mut self, other: Population<T, RatedSorted>, n: usize) {
        self.individuals.extend(other.individuals.into_iter().take(n));
    }

    /// Append all individuals of population `other`.
    pub fn append<X: IsRated>(&mut self, other: Population<T, X>) {
        self.individuals.extend(other.individuals.into_iter());
    }

    // We want to create a population of approximately `pop_size`. We do not care if there are
    // slightly more or slightly less individuals than that.
    //
    // 1. sort whole popluation into niches
    // 2. calculate the number of offspring for each niche.
    // 3. each niche produces offspring.
    //    - sort each niche according to the fitness value.
    //    - r% of the best genomes produce offspring
    //    - determine the elitist size. those are always copied into the new generation.
    //
    pub fn reproduce<C, M, R>(self,
                              pop_size: usize,
                              // how many of the best individuals of a niche are copied as-is into the
                              // new population?
                              elite_percentage: Closed01<f64>,
                              // how many of the best individuals of a niche are selected for
                              // reproduction?
                              selection_percentage: Closed01<f64>,
                              compatibility_threshold: f64,
                              compatibility: &C,
                              mate: &mut M,
                              rng: &mut R)
                              -> (Population<T, Rated>, Population<T, Unrated>)
        where R: Rng,
              C: Distance<T>,
              M: Mate<T>
    {
        debug!("population size: {}", self.len());

        assert!(elite_percentage <= selection_percentage); // XXX
        assert!(self.len() > 0);
        let niches = self.partition(rng, compatibility_threshold, compatibility);
        let num_niches = niches.len();

        debug!("number of niches: {}", num_niches);

        assert!(num_niches > 0);
        let total_mean = niches.total_mean();
        assert!(total_mean.get() >= 0.0);

        let mut new_unrated_population: Population<T, Unrated> = Population::new();
        let mut new_rated_population: Population<T, Rated> = Population::new();

        for niche in niches.into_iter() {
            let percentage_of_population: f64 = if total_mean.get() == 0.0 {
                // all individuals have a fitness of 0.0.
                // we will equally allow each niche to procude offspring.
                1.0 / (num_niches as f64)
            } else {
                (niche.mean_fitness() / total_mean).get()
            };

            // calculate new size of niche, and size of elites, and selection size.
            assert!(percentage_of_population >= 0.0 && percentage_of_population <= 1.0);

            let niche_size = pop_size as f64 * percentage_of_population;

            // number of elitary individuals to copy from the old niche generation into the new.
            let elite_size =
                cmp::max(1,
                         probabilistic_round(niche_size * elite_percentage.get(), rng) as usize);

            // number of offspring to produce.
            let offspring_size = probabilistic_round(niche_size * (1.0 - elite_percentage.get()),
                                                     rng) as usize;

            // number of the best individuals to use for mating.
            let select_size = probabilistic_round(niche_size *
                                                  selection_percentage.get(),
                                                  rng) as usize;

            let sorted_niche = niche.sort();

            let select_size = cmp::min(select_size, sorted_niche.len());

            // at first produce `offspring_size` individuals from the top `select_size`
            // individuals.
            if select_size > 0 {
                for _ in 0..offspring_size {
                    let offspring = sorted_niche.create_single_offspring(select_size, mate, rng);
                    new_unrated_population.add_genome(Box::new(offspring));
                }
            }

            // now copy the elites
            new_rated_population.merge(sorted_niche, elite_size);
        }

        return (new_rated_population, new_unrated_population);
    }

    /// Partition the whole population into species (niches)

    pub fn partition<R, C>(self,
                           rng: &mut R,
                           compatibility_threshold: f64,
                           compatibility: &C)
                           -> Niches<T>
        where R: Rng,
              C: Distance<T>
    {
        let mut niches = Niches::new();

        for ind in self.individuals.into_iter() {
            niches.add_individual(ind, compatibility_threshold, compatibility, rng);
        }

        niches
    }
}

pub struct Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    // anticipated population size
    pub pop_size: usize,
    // how many of the best individuals of a niche are copied as-is into the
    // new population?
    pub elite_percentage: Closed01<f64>,
    // how many of the best individuals of a niche are selected for
    // reproduction?
    pub selection_percentage: Closed01<f64>,
    pub compatibility_threshold: f64,
    pub compatibility: &'a C,
    pub mate: &'a mut M,
    pub fitness: &'a F,
    pub _marker: PhantomData<T>,
}

impl<'a, T, C, M, F> Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    pub fn run<R, G>(&mut self,
                     initial_pop: Population<T, Unrated>,
                     goal_condition: &G,
                     rng: &mut R)
                     -> (usize, Population<T, Rated>)
        where R: Rng,
              G: Fn(usize, &Population<T, Rated>) -> bool
    {
        let mut iteration: usize = 0;
        let mut current_rated_pop = initial_pop.rate_par(self.fitness);

        while !goal_condition(iteration, &current_rated_pop) {
            let (new_rated, new_unrated) =
                current_rated_pop.reproduce(self.pop_size,
                                            self.elite_percentage,
                                            self.selection_percentage,
                                            self.compatibility_threshold,
                                            self.compatibility,
                                            self.mate,
                                            rng);

            current_rated_pop = new_rated;
            current_rated_pop.append(new_unrated.rate_par(self.fitness));
            iteration += 1;
        }

        return (iteration, current_rated_pop);
    }
}
<commit_msg>Use Option<Fitness> instead of a Zero Fitness. Makes it more safe at runtime.<commit_after>use super::fitness::Fitness;
use super::traits::{Genotype, Distance};
use rand::Rng;
use closed01::Closed01;
use super::traits::Mate;
use std::marker::PhantomData;
use std::num::Zero;
use std::fmt::Debug;
use std::cmp;
use std::mem;
use rayon::par_iter::*;
use super::prob::probabilistic_round;

#[derive(Debug)]
pub struct Individual<T: Debug + Genotype> {
    fitness: Option<Fitness>,
    genome: Box<T>,
}

impl<T: Debug + Genotype> Individual<T> {
    pub fn fitness(&self) -> Fitness {
        self.fitness.unwrap()
    }

    pub fn genome(&self) -> &T {
        &self.genome
    }
}

pub trait Rating { }

pub trait IsRated : Rating { }
pub trait IsRatedSorted : IsRated { }

#[derive(Debug)]
pub struct Unrated;

#[derive(Debug)]
pub struct Rated;

#[derive(Debug)]
pub struct RatedSorted;

impl Rating for Unrated {}
impl Rating for Rated {}
impl Rating for RatedSorted {}

impl IsRated for Rated {}
impl IsRated for RatedSorted {}
impl IsRatedSorted for RatedSorted {}

#[derive(Debug)]
pub struct Population<T: Genotype + Debug, R: Rating> {
    individuals: Vec<Individual<T>>,
    _marker: PhantomData<R>,
}

#[derive(Debug)]
pub struct Niches<T: Genotype + Debug> {
    total_individuals: usize,
    niches: Vec<Population<T, Rated>>,
}

impl<T: Genotype + Debug> Niches<T> {
    pub fn new() -> Self {
        Niches {
            total_individuals: 0,
            niches: Vec::new(),
        }
    }

    pub fn into_iter(self) -> ::std::vec::IntoIter<Population<T, Rated>> {
        self.niches.into_iter()
    }

    /// The sum of all "mean fitnesses" of all niches.

    fn total_mean(&self) -> Fitness {
        self.niches.iter().map(|ind| ind.mean_fitness()).sum()
    }

    /// Number of niches

    pub fn len(&self) -> usize {
        self.niches.len()
    }

    /// Add an individual to one of the niches.

    pub fn add_individual<R, C>(&mut self,
                                ind: Individual<T>,
                                compatibility_threshold: f64,
                                compatibility: &C,
                                rng: &mut R)
        where R: Rng,
              C: Distance<T>
    {
        self.total_individuals += 1;

        for niche in self.niches.iter_mut() {
            // Is this genome compatible with this niche? Test against a random genome.
            let compatible = match rng.choose(&niche.individuals) {
                Some(probe) => {
                    compatibility.distance(&probe.genome, &ind.genome) < compatibility_threshold
                }
                // If a niche is empty, a genome always is compatible (note that a niche can't be empyt)
                None => true,
            };
            if compatible {
                niche.add_individual(ind);
                return;
            }
        }
        // if no compatible niche was found, create a new niche containing this genome.
        let mut new_niche = Population::new();
        new_niche.add_individual(ind);
        self.niches.push(new_niche);
    }
}

impl<T: Genotype + Debug, R: Rating> Population<T, R> {
    pub fn new() -> Population<T, R> {
        Population {
            individuals: Vec::new(),
            _marker: PhantomData,
        }
    }

    pub fn new_from_vec(v: Vec<Individual<T>>) -> Population<T, R> {
        Population {
            individuals: v,
            _marker: PhantomData,
        }
    }

    pub fn len(&self) -> usize {
        self.individuals.len()
    }
}

impl<T: Genotype + Debug> Population<T, Unrated> {
    pub fn add_genome(&mut self, genome: Box<T>) {
        self.individuals.push(Individual {
            fitness: None,
            genome: genome,
        });
    }

    pub fn rate_seq<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Fn(&T) -> Fitness
    {
        for ind in self.individuals.iter_mut() {
            let fitness = f(&ind.genome);
            ind.fitness = Some(fitness);
        }
        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }

    pub fn rate_par<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Sync + Fn(&T) -> Fitness
    {
        self.individuals.par_iter_mut().for_each(|ind| {
            let fitness = f(&ind.genome);
            ind.fitness = Some(fitness);
        });

        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }
}

impl<T: Genotype + Debug, R: IsRated> Population<T, R> {
    fn mean_fitness(&self) -> Fitness {
        let sum: Fitness = self.individuals.iter().map(|ind| ind.fitness()).sum();
        sum / Fitness::new(self.len() as f64)
    }
}

impl<T: Genotype + Debug> Population<T, RatedSorted> {

    pub fn into_iter(self) -> ::std::vec::IntoIter<Individual<T>> {
        self.individuals.into_iter()
    }

    pub fn best_individual(&self) -> Option<&Individual<T>> {
        self.individuals.first()
    }

    // Return true if genome at position `i` is fitter that `j`
    //
    // In a sorted population, the individual with the lower index
    // has a better fitness.

    // #[inline]
    // fn is_fitter(&self, i: usize, j: usize) -> bool {
    //    i < j
    // }

    /// Create a single offspring Genome by selecting random parents
    /// from the best `select_size` individuals of the populations.

    fn create_single_offspring<R, M>(&self, select_size: usize, mate: &mut M, rng: &mut R) -> T
        where R: Rng,
              M: Mate<T>
    {
        assert!(select_size > 0 && select_size <= self.len());

        // We do not need tournament selection here as our population is sorted.
        // We simply determine two individuals out of `select_size`.

        let mut parent1 = rng.gen_range(0, select_size);
        let mut parent2 = rng.gen_range(0, select_size);

        // try to find a parent2 != parent1. retry three times.
        for _ in 0..3 {
            if parent2 != parent1 {
                break;
            }
            parent2 = rng.gen_range(0, select_size);
        }

        // `mate` assumes that the first parent performs better.
        if parent1 > parent2 {
            mem::swap(&mut parent1, &mut parent2);
        }

        debug_assert!(parent1 <= parent2);

        mate.mate(&self.individuals[parent1].genome,
                  &self.individuals[parent2].genome,
                  parent1 == parent2,
                  rng)
    }
}

impl<T: Genotype + Debug> Population<T, Rated> {
    pub fn add_individual(&mut self, ind: Individual<T>) {
        self.individuals.push(ind);
    }

    // Return true if genome at position `i` is fitter that `j`
    // fn is_fitter(&self, i: usize, j: usize) -> bool {
    //    self.individuals[i].fitness > self.individuals[j].fitness
    // }

    // higher value of fitness means that the individual is fitter.
    pub fn sort(mut self) -> Population<T, RatedSorted> {
        (&mut self.individuals).sort_by(|a, b| a.fitness().cmp(&b.fitness()).reverse());
        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }

    pub fn best_individual(&self) -> Option<&Individual<T>> {
        self.individuals.iter().max_by_key(|ind| ind.fitness())
    }

    /// Merge `self` with the first `n` individuals from population `other`.
    pub fn merge(&mut self, other: Population<T, RatedSorted>, n: usize) {
        self.individuals.extend(other.individuals.into_iter().take(n));
    }

    /// Append all individuals of population `other`.
    pub fn append<X: IsRated>(&mut self, other: Population<T, X>) {
        self.individuals.extend(other.individuals.into_iter());
    }

    // We want to create a population of approximately `pop_size`. We do not care if there are
    // slightly more or slightly less individuals than that.
    //
    // 1. sort whole popluation into niches
    // 2. calculate the number of offspring for each niche.
    // 3. each niche produces offspring.
    //    - sort each niche according to the fitness value.
    //    - r% of the best genomes produce offspring
    //    - determine the elitist size. those are always copied into the new generation.
    //
    pub fn reproduce<C, M, R>(self,
                              pop_size: usize,
                              // how many of the best individuals of a niche are copied as-is into the
                              // new population?
                              elite_percentage: Closed01<f64>,
                              // how many of the best individuals of a niche are selected for
                              // reproduction?
                              selection_percentage: Closed01<f64>,
                              compatibility_threshold: f64,
                              compatibility: &C,
                              mate: &mut M,
                              rng: &mut R)
                              -> (Population<T, Rated>, Population<T, Unrated>)
        where R: Rng,
              C: Distance<T>,
              M: Mate<T>
    {
        debug!("population size: {}", self.len());

        assert!(elite_percentage <= selection_percentage); // XXX
        assert!(self.len() > 0);
        let niches = self.partition(rng, compatibility_threshold, compatibility);
        let num_niches = niches.len();

        debug!("number of niches: {}", num_niches);

        assert!(num_niches > 0);
        let total_mean = niches.total_mean();
        assert!(total_mean.get() >= 0.0);

        let mut new_unrated_population: Population<T, Unrated> = Population::new();
        let mut new_rated_population: Population<T, Rated> = Population::new();

        for niche in niches.into_iter() {
            let percentage_of_population: f64 = if total_mean.get() == 0.0 {
                // all individuals have a fitness of 0.0.
                // we will equally allow each niche to procude offspring.
                1.0 / (num_niches as f64)
            } else {
                (niche.mean_fitness() / total_mean).get()
            };

            // calculate new size of niche, and size of elites, and selection size.
            assert!(percentage_of_population >= 0.0 && percentage_of_population <= 1.0);

            let niche_size = pop_size as f64 * percentage_of_population;

            // number of elitary individuals to copy from the old niche generation into the new.
            let elite_size =
                cmp::max(1,
                         probabilistic_round(niche_size * elite_percentage.get(), rng) as usize);

            // number of offspring to produce.
            let offspring_size = probabilistic_round(niche_size * (1.0 - elite_percentage.get()),
                                                     rng) as usize;

            // number of the best individuals to use for mating.
            let select_size = probabilistic_round(niche_size *
                                                  selection_percentage.get(),
                                                  rng) as usize;

            let sorted_niche = niche.sort();

            let select_size = cmp::min(select_size, sorted_niche.len());

            // at first produce `offspring_size` individuals from the top `select_size`
            // individuals.
            if select_size > 0 {
                for _ in 0..offspring_size {
                    let offspring = sorted_niche.create_single_offspring(select_size, mate, rng);
                    new_unrated_population.add_genome(Box::new(offspring));
                }
            }

            // now copy the elites
            new_rated_population.merge(sorted_niche, elite_size);
        }

        return (new_rated_population, new_unrated_population);
    }

    /// Partition the whole population into species (niches)

    pub fn partition<R, C>(self,
                           rng: &mut R,
                           compatibility_threshold: f64,
                           compatibility: &C)
                           -> Niches<T>
        where R: Rng,
              C: Distance<T>
    {
        let mut niches = Niches::new();

        for ind in self.individuals.into_iter() {
            niches.add_individual(ind, compatibility_threshold, compatibility, rng);
        }

        niches
    }
}

pub struct Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    // anticipated population size
    pub pop_size: usize,
    // how many of the best individuals of a niche are copied as-is into the
    // new population?
    pub elite_percentage: Closed01<f64>,
    // how many of the best individuals of a niche are selected for
    // reproduction?
    pub selection_percentage: Closed01<f64>,
    pub compatibility_threshold: f64,
    pub compatibility: &'a C,
    pub mate: &'a mut M,
    pub fitness: &'a F,
    pub _marker: PhantomData<T>,
}

impl<'a, T, C, M, F> Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    pub fn run<R, G>(&mut self,
                     initial_pop: Population<T, Unrated>,
                     goal_condition: &G,
                     rng: &mut R)
                     -> (usize, Population<T, Rated>)
        where R: Rng,
              G: Fn(usize, &Population<T, Rated>) -> bool
    {
        let mut iteration: usize = 0;
        let mut current_rated_pop = initial_pop.rate_par(self.fitness);

        while !goal_condition(iteration, &current_rated_pop) {
            let (new_rated, new_unrated) =
                current_rated_pop.reproduce(self.pop_size,
                                            self.elite_percentage,
                                            self.selection_percentage,
                                            self.compatibility_threshold,
                                            self.compatibility,
                                            self.mate,
                                            rng);

            current_rated_pop = new_rated;
            current_rated_pop.append(new_unrated.rate_par(self.fitness));
            iteration += 1;
        }

        return (iteration, current_rated_pop);
    }
}
<|endoftext|>