packages feed

telescope-0.3.0: test/Test/Fits/MegaHeaderSpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Test.Fits.MegaHeaderSpec where

import Control.Monad (forM_)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as C8
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Skeletest
import System.IO
import Telescope.Data.Axes
import Telescope.Fits.BitPix
import Telescope.Fits.DataArray (Dimensions (..))
import Telescope.Fits.Encoding.MegaHeader
import Telescope.Fits.Header (Header (..), HeaderRecord (..), KeywordRecord (..), LogicalConstant (..), Value (..), keywords, lookupKeyword)
import Text.Megaparsec qualified as M
import Text.Megaparsec.Char qualified as M
import Prelude hiding (lookup)


spec :: Spec
spec = do
  basicParsing
  keywordValueLines
  comments
  continue
  fullRecord
  fullRecordLine
  headerMap
  requiredHeaders
  sampleNSOHeaders
  sampleHubbleHeaders
  sampleJWSTHeaders


parse :: Parser a -> ByteString -> IO a
parse p inp =
  case M.parse p "Test" inp of
    Left e -> fail $ showParseError e
    Right v -> pure v


flattenKeywords :: [ByteString] -> ByteString
flattenKeywords ts = mconcat (map pad ts) <> "END"


pad :: ByteString -> ByteString
pad m =
  let n = 80 - BS.length m
   in m <> C8.replicate n ' '


basicParsing :: Spec
basicParsing = describe "Basic Parsing" $ do
  it "should parse a string" $ do
    res <- parse parseStringValue "'hello there' "
    res `shouldBe` "hello there"

  it "should parse a number value" $ do
    res <- parse parseValue "42 "
    res `shouldBe` Integer 42

  it "should parse a keyword" $ do
    res <- parse parseKeyword "WS_TEMP ="
    res `shouldBe` "WS_TEMP"

  it "should handle keyword symbols" $ do
    res <- parse parseKeyword "OBSGEO-X=   -5466045.256954942 / [m]"
    res `shouldBe` "OBSGEO-X"


keywordValueLines :: Spec
keywordValueLines = describe "parse keyword=value" $ do
  it "should parse an integer" $ do
    res <- parse parseKeywordValue "KEY=42 "
    res `shouldBe` ("KEY", Integer 42)

  it "should parse a string" $ do
    res <- parse parseKeywordValue "KEY='value'"
    res `shouldBe` ("KEY", String "value")

  it "should absorb spaces" $ do
    res <- parse parseKeywordValue "KEY   = 'value'   "
    res `shouldBe` ("KEY", String "value")

  it "should parse a float" $ do
    res <- parse parseKeywordValue "KEY =   44.88 "
    res `shouldBe` ("KEY", Float 44.88)

  it "should parse a negative number" $ do
    res <- parse parseKeywordValue "KEY = -44.88"
    res `shouldBe` ("KEY", Float (-44.88))

  it "should parse a logical constant" $ do
    res <- parse parseKeywordValue "KEYT=     T "
    res `shouldBe` ("KEYT", Logic T)

    res2 <- parse parseKeywordValue "KEYF=     F "
    res2 `shouldBe` ("KEYF", Logic F)

  it "should ignore comments" $ do
    res <- parse parseKeywordValue "SIMPLE  =                    T / conforms to FITS standard"
    res `shouldBe` ("SIMPLE", Logic T)

  it "should strip trailing spaces from strings" $ do
    res <- parse parseKeywordValue "INSTRUME= 'VISP    '"
    res `shouldBe` ("INSTRUME", String "VISP")


fullRecord :: Spec
fullRecord = describe "parseKeywordRecord" $ do
  it "should parse an 80 character record" $ do
    res <- parse parseKeywordRecord (flattenKeywords ["KEYWORD = 12345"])
    res `shouldBe` KeywordRecord "KEYWORD" (Integer 12345) Nothing

  it "should parse an a record and comment" $ do
    res <- parse parseKeywordRecord (flattenKeywords ["KEYWORD = 12345 / this is a comment"])
    res `shouldBe` KeywordRecord "KEYWORD" (Integer 12345) (Just "this is a comment")

  it "should parse a record, comment, followed by next keyword" $ do
    res <- parse parseKeywordRecord $ flattenKeywords ["SIMPLE  =                    T / conforms to FITS standard"]
    res `shouldBe` KeywordRecord "SIMPLE" (Logic T) (Just "conforms to FITS standard")

  it "should handle keyword symbols" $ do
    res <- parse parseKeywordRecord $ flattenKeywords ["OBSGEO-X=   -5466045.256954942 / [m]"]
    res `shouldBe` KeywordRecord "OBSGEO-X" (Float (-5466045.256954942)) (Just "[m]")

  it "should handle extension" $ do
    res <- parse parseKeywordRecord $ flattenKeywords ["XTENSION= 'IMAGE   '"]
    res `shouldBe` KeywordRecord "XTENSION" (String "IMAGE") Nothing

  it "should not consume blank lines" $ do
    let inp = flattenKeywords ["SIMPLE  =                    T", " "]
    res <- flip parse inp $ do
      kv <- parseKeywordRecord
      _ <- parseLineBlank
      pure kv
    res `shouldBe` KeywordRecord "SIMPLE" (Logic T) Nothing


fullRecordLine :: Spec
fullRecordLine = describe "parseRecordLine" $ do
  it "should parse a normal line" $ do
    res <- parse parseRecordLine "NAXIS1  =                  100 / [pix]                                          END"
    res `shouldBe` Keyword (KeywordRecord "NAXIS1" (Integer 100) (Just "[pix]"))

  it "should parse a comment line" $ do
    res <- parse parseRecordLine "COMMENT ------------------------------ Telescope -------------------------------END"
    res `shouldBe` Comment "------------------------------ Telescope -------------------------------"

  it "should parse a blank line" $ do
    res <- parse parseRecordLine $ flattenKeywords [" "]
    res `shouldBe` BlankLine


comments :: Spec
comments = do
  describe "Full-line comments" $ withMarkers ["focus"] $ do
    it "should parse full-line comments" $ do
      res <- parse parseLineComment $ flattenKeywords ["COMMENT --------------------------- VISP Instrument ----------------------------"]
      res `shouldBe` "--------------------------- VISP Instrument ----------------------------"

    it "should parse comments with text" $ do
      res <- parse parseLineComment $ flattenKeywords ["COMMENT  Keys describing the pointing and operation of the telescope. Including "]
      res `shouldBe` " Keys describing the pointing and operation of the telescope. Including "

    it "should parse blank comments" $ do
      res <- parse parseLineComment $ flattenKeywords ["COMMENT                                                                         "]
      res `shouldBe` "                                                                        "

    it "should parse silly JWST blank-line comments" $ do
      res <- parse parseLineCommentSpaces $ flattenKeywords ["        Association information                                                "]
      res `shouldBe` "Association information"

    it "should parse silly JWST blank-line comments as a header" $ do
      let hs =
            flattenKeywords
              [ "                                                                                "
              , "        Association information                                                 "
              , "                                                                                "
              ]

      h <- parse parseHeader hs
      length h.records `shouldBe` 3

    it "fails space-line comments on blank lines" $ do
      res <- parse parseRecordLine (flattenKeywords ["                                                                                "])
      res `shouldBe` BlankLine

  describe "inline comments" $ do
    it "should parse comment" $ do
      res <- parse (parseInlineComment 0) $ " / Telescope" <> C8.replicate 68 ' '
      res `shouldBe` "Telescope"

  describe "parse to line end" $ do
    it "should parse comment line end" $ do
      res <- parse (parseLineEnd 0) $ " / Telescope" <> C8.replicate 68 ' '
      res `shouldBe` Just "Telescope"

    it "should ignore blanks" $ do
      res <- parse (parseLineEnd 0) $ C8.replicate 80 ' '
      res `shouldBe` Nothing

    it "should parse comment after blanks" $ do
      res <- parse (parseLineEnd 0) $ "          / comment " <> C8.replicate 60 ' '
      res `shouldBe` Just "comment"

  describe "withComments" $ do
    it "should parse a number, ignoring" $ do
      res <- parse (withComments parseValue) $ "12345 / Woot" <> C8.replicate 68 ' '
      res `shouldBe` (Integer 12345, Just "Woot")

    it "should parse no comment" $ do
      res <- parse (withComments parseValue) $ "12345       " <> C8.replicate 68 ' '
      res `shouldBe` (Integer 12345, Nothing)

    it "should ignore comments on bintable" $ do
      -- parse (withComments parseValue) $ "12345       " <> C8.replicate 68 ' '
      let inp = flattenKeywords ["XTENSION= 'BINTABLE'           / binary table extension"]
      (_, mc) <- flip parse inp $ do
        withComments $ M.string' "XTENSION= 'BINTABLE'"
      mc `shouldBe` Just "binary table extension"


continue :: Spec
continue = describe "Continue Keyword" $ withMarkers ["continue"] $ do
  it "parses a normal string" $ do
    let input = "'normal string  '     END"
    t <- parse parseStringContinue input
    t `shouldBe` "normal string"

  it "parses a CONTINUE with space after" $ do
    let input = "'has some sp&'CONTINUE  'ace'                                                                  "
    t <- parse parseStringContinue input
    t `shouldBe` "has some space"

  it "parses a CONTINUE with comment" $ do
    let input = "'has a comm'CONTINUE  'ent'    / hello world                                                 "
    t <- parse parseStringContinue input
    t `shouldBe` "has a comment"

  it "parses a short CONTINUE" $ do
    let input = "'stops before the e'  CONTINUE  'nd'                                                                  "
    t <- parse parseStringContinue input
    t `shouldBe` "stops before the end"

  it "combines continue into previous keyword" $ do
    let hs =
          [ "CAL_URL = 'https://docs.dkist.nso.edu/projects/visp/en/v2.0.1/l0_to_l1_visp.ht&'"
          , "CONTINUE  'ml'                                                                  "
          ]

    h <- parse parseHeader $ flattenKeywords hs
    lookupKeyword "CAL_URL" h `shouldBe` Just (String "https://docs.dkist.nso.edu/projects/visp/en/v2.0.1/l0_to_l1_visp.html")

  it "parses funny hubble full-line comments" $ do
    -- Hubble has a comment KEYWORD at the end, which doesn't go fully to the end of the line
    let input =
          flattenKeywords
            [ "COMMENT = 'If you have used HLA files for your research, please include the &'  "
            , "CONTINUE  'following acknowledgment:&'                                          "
            ]

    h <- parse parseHeader input
    h.records `shouldBe` [Keyword $ KeywordRecord "COMMENT" (String "If you have used HLA files for your research, please include the following acknowledgment:&") Nothing]


headerMap :: Spec
headerMap = describe "full header" $ do
  it "should parse single header" $ do
    h <- parse parseHeader $ flattenKeywords ["KEY1='value'"]
    length (keywords h) `shouldBe` 1
    lookupKeyword "KEY1" h `shouldBe` Just (String "value")

  it "should parse multiple headers " $ do
    res <- parse parseHeader $ flattenKeywords ["KEY1='value'", "KEY2=  23"]
    length (keywords res) `shouldBe` 2
    lookupKeyword "KEY2" res `shouldBe` Just (Integer 23)

  it "should ignore comments" $ do
    res <- parse parseHeader $ flattenKeywords ["KEY1='value' / this is a comment"]
    length (keywords res) `shouldBe` 1
    lookupKeyword "KEY1" res `shouldBe` Just (String "value")

  it "should handle xtension" $ do
    res <- parse parseHeader $ flattenKeywords ["XTENSION= 'IMAGE   '"]
    length (keywords res) `shouldBe` 1
    lookupKeyword "XTENSION" res `shouldBe` Just (String "IMAGE")

  it "should parse blank line before keyword" $ do
    res <- parse parseHeader $ flattenKeywords [" ", "KEY2    = 22"]
    res.records `shouldBe` [BlankLine, Keyword (KeywordRecord "KEY2" (Integer 22) Nothing)]

  it "should parse line after keyword" $ do
    res <- parse parseHeader $ flattenKeywords ["KEY1    = 11", " ", "KEY2    = 22"]
    res.records `shouldBe` [Keyword (KeywordRecord "KEY1" (Integer 11) Nothing), BlankLine, Keyword (KeywordRecord "KEY2" (Integer 22) Nothing)]


requiredHeaders :: Spec
requiredHeaders =
  describe "required headers" $ do
    it "should parse bitpix" $ do
      res <- parse parseBitPix $ flattenKeywords ["BITPIX = 16"]
      res `shouldBe` BPInt16

    it "should parse NAXES" $ do
      res <- parse parseNaxes $ flattenKeywords ["NAXIS   =3", "NAXIS1  =1", "NAXIS2  =2", "NAXIS3  =3"]
      res `shouldBe` Axes [1, 2, 3]

    it "should parse size" $ do
      res <- parse parseDimensions $ flattenKeywords ["BITPIX = -32", "NAXIS=2", "NAXIS1=10", "NAXIS2=20"]
      res.bitpix `shouldBe` BPFloat
      res.axes `shouldBe` Axes [10, 20]

    it "should parse data headers with data" $ do
      let fakeData = "1234" -- Related to NAXIS!
      h <- parse parseHeader $ flattenKeywords ["SIMPLE = T", "BITPIX = 8", "NAXIS=2", "NAXIS1=2", "NAXIS2=2", "TEST='hi'"] <> fakeData
      length (keywords h) `shouldBe` 6
      lookupKeyword "NAXIS" h `shouldBe` Just (Integer 2)


sampleNSOHeaders :: Spec
sampleNSOHeaders = do
  describe "NSO Stripped Headers" $ do
    it "should parse comment block" $ do
      let hs =
            [ "DATASUM = '550335088'          / data unit checksum updated 2023-04-22T04:10:59 "
            , "   "
            , "COMMENT ------------------------------ Telescope -------------------------------"
            , "COMMENT  Keys describing the pointing and operation of the telescope. Including "
            , "COMMENT     the FITS WCS keys describing the world coordinates of the array.    "
            ]
      h <- parse parseHeader $ flattenKeywords hs
      length (keywords h) `shouldBe` 1

    describe "sample header file" $ do
      it "should parse all keywords individually" $ do
        SampleNSO ts <- getFixture
        forM_ (zip [1 :: Int ..] ts) $ \(_, t) -> do
          _ <- parse parseRecordLine $ flattenKeywords [TE.encodeUtf8 t]
          pure ()

      it "should parse bin table keywords" $ do
        DKISTHeaders bs <- getFixture
        (sz, pc) <- parse parseBinTableKeywords (mconcat $ C8.lines bs)
        sz.axes `shouldBe` Axes [32, 998]
        pc `shouldBe` 95968


sampleHubbleHeaders :: Spec
sampleHubbleHeaders = do
  describe "Hubble Headers" $ do
    it "should parse weird comment lines" $ do
      let hs =
            [ "SIMPLE  =                    T / conforms to FITS standard                      "
            , "BITPIX  =                    8 / array data type                                "
            , "NAXIS   =                    0 / number of array dimensions                     "
            , "EXTEND  =                    T                                                  "
            , "DATE    = '2009-11-10'         / date this file was written (yyyy-mm-dd)        "
            , "FILETYPE= 'SCI      '          / type of data found in data file                "
            , "                                                                                "
            , "TELESCOP= 'HST'                / telescope used to acquire data                 "
            , "INSTRUME= 'WFPC2 '             / identifier for instrument used to acquire data "
            , "EQUINOX =               2000.0 / equinox of celestial coord. system             "
            , "                                                                                "
            , "              / WFPC-II DATA DESCRIPTOR KEYWORDS                                "
            , "                                                                                "
            , "ROOTNAME= 'hst_8599_53_wfpc2_f814w_wf' / rootname of the observation set        "
            ]

      h <- parse parseHeader $ flattenKeywords hs
      length (keywords h) `shouldBe` 10

    it "unknown error reproduce" $ do
      let hs =
            [ "EQUINOX =              '2000.0'                                                 "
            , "                                                                                "
            , "              / WFPC-II DATA DESCRIPTOR KEYWORDS                                "
            , "ROOTNAME= 'hst_8599_53_wfpc2_f814w_wf' / rootname of the observation set        "
            ]

      h <- parse parseHeader $ flattenKeywords hs
      length (keywords h) `shouldBe` 2

    it "should parse hubble headers" $ do
      let hs1 =
            flattenKeywords
              [ "SIMPLE  =                    T / conforms to FITS standard                      "
              , "BITPIX  =                    8 / array data type                                "
              , "NAXIS   =                    0 / number of array dimensions                     "
              , "EXTEND  =                    T                                                  "
              , "DATE    = '2009-11-10'         / date this file was written (yyyy-mm-dd)        "
              ]
      HubbleHeaders bs <- getFixture
      let input = mconcat $ C8.lines bs
      BS.take 400 input <> "END" `shouldBe` hs1

      h <- parse parseHeader $ BS.take (601 * 80) input <> "END"
      lookupKeyword "DATE-OBS" h `shouldBe` Just (String "2001-04-07")
      lookupKeyword "FILTROT" h `shouldBe` Just (Float 0)


sampleJWSTHeaders :: Spec
sampleJWSTHeaders = do
  describe "JWST Headers" $ withMarkers ["focus"] $ do
    it "should parse jwst headers" $ do
      JWSTHeaders bs <- getFixture
      let input = mconcat $ C8.lines bs
      h <- parse parseHeader input
      length (keywords h) `shouldBe` 214


newtype DKISTHeaders = DKISTHeaders BS.ByteString
instance Fixture DKISTHeaders where
  fixtureAction = do
    bs <- BS.readFile "./samples/nso_dkist_headers.txt"
    pure $ noCleanup $ DKISTHeaders bs


newtype HubbleHeaders = HubbleHeaders BS.ByteString
instance Fixture HubbleHeaders where
  fixtureAction = do
    bs <- BS.readFile "./samples/hubble_headers.txt"
    pure $ noCleanup $ HubbleHeaders bs


newtype JWSTHeaders = JWSTHeaders BS.ByteString
instance Fixture JWSTHeaders where
  fixtureAction = do
    bs <- BS.readFile "./samples/jwst_headers.txt"
    pure $ noCleanup $ JWSTHeaders bs


newtype SampleNSO = SampleNSO [Text]
instance Fixture SampleNSO where
  fixtureAction = do
    DKISTHeaders bs <- getFixture
    pure $ noCleanup $ SampleNSO $ filter (not . ignore) $ T.lines $ TE.decodeUtf8 bs
   where
    ignore t = T.isPrefixOf "CONTINUE" t || T.isPrefixOf "END" t