Full Code of flannelhead/blackstar for AI

Repository: flannelhead/blackstar
Branch: master
Commit: 40f1f1965292
Files: 26
Total size: 43.3 KB

Directory structure:
gitextract_5lpryipi/

├── .gitignore
├── LICENSE
├── README.md
├── Setup.hs
├── animations/
│   └── default-ani.yaml
├── app/
│   ├── Animate.hs
│   ├── GenerateTree.hs
│   └── Main.hs
├── blackstar.cabal
├── scenes/
│   ├── closeup.yaml
│   ├── default-aa.yaml
│   ├── default.yaml
│   ├── fartheraway.yaml
│   ├── lensing-disk.yaml
│   ├── lensing.yaml
│   ├── wideangle-disk.yaml
│   ├── wideangle.yaml
│   └── wideangle1.yaml
├── scripts/
│   └── ffmpeg-animate
├── src/
│   ├── Animation.hs
│   ├── ConfigFile.hs
│   ├── ImageFilters.hs
│   ├── Raytracer.hs
│   ├── StarMap.hs
│   └── Util.hs
└── stack.yaml

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
.stack-work
blackstar.prof
out.png
texture.jpg
BSC5
SAO.pc
PPM
PPMra
renders
bloomed.png
output
*.png
!example.png
frames*
*.mkv
dist*
.ghc.*



================================================
FILE: LICENSE
================================================
Copyright Sakari Kapanen (c) 2016

All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.

    * Redistributions in binary form must reproduce the above
      copyright notice, this list of conditions and the following
      disclaimer in the documentation and/or other materials provided
      with the distribution.

    * Neither the name of Author name here nor the names of other
      contributors may be used to endorse or promote products derived
      from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


================================================
FILE: README.md
================================================
# Blackstar
A black hole ray tracer written in Haskell. There's [an article](https://flannelhead.github.io/projects/blackstar.html) about this on my homepage. I've also written a [theoretical writeup](https://flannelhead.github.io/posts/2016-03-06-photons-and-black-holes.html) on Schwarzschild geodesics.

![An example image](https://raw.githubusercontent.com/flannelhead/blackstar/master/example.png)

## Features
* Fast, parallel ray tracing
* Rendering [Schwarzschild](https://en.wikipedia.org/wiki/Schwarzschild_metric) black holes
* Rendering accretion disks
* Drawing the celestial sphere using a star catalogue
* Bloom effect
* Antialiasing by 4x supersampling for smoother images
* Easy, YAML based configuration
* A simple CLI
* Batch mode and sequence generator for creating animations

## What about the name?
It is a tribute to David Bowie, referring to his last album.

## Building
Use [`stack`](http://docs.haskellstack.org/en/stable/README/) to build this. First clone the repo, then run `stack build` and follow the instructions given by `stack`. You should be able to build `blackstar` on any platform where you can install `stack`.

You will have to build the star lookup tree first. Download the [PPM star catalog](http://tdc-www.harvard.edu/software/catalogs/ppm.html) [this archive](http://tdc-www.harvard.edu/software/catalogs/ppm.tar.gz) and extract the file `PPM` to the root folder of this project. Then run `stack exec generate-tree PPM stars.kdt` and the tree will be generated and saved to the file `stars.kdt`.

### Speeding it up with LLVM
When doing large or batch renders, it is recommended to build `blackstar` using GHC's LLVM backend. GHC produces LLVM bytecode and LLVM produces fast native code from GHC's output. In my tests I've noticed ~1.5x speedups.

The LLVM backend isn't used by default since one needs to install (and usually build) a specific version of LLVM separately. Moreover, the build time is significantly higher with LLVM, so one doesn't definitely want to use it while hacking on the code.

To successfully build with LLVM, you need to:

* Download and [build](http://llvm.org/docs/GettingStarted.html#getting-started-quickly-a-summary) [LLVM 6.0.1](http://llvm.org/releases/download.html#6.0.1). You can skip the Clang parts. After the build, you should make sure the tools `llc` and `opt` are found in your `PATH`. Notice that these aren't included in the prebuilt LLVM binaries, that's why you'll need to build it.
* Build `blackstar` with `stack build --ghc-options -fllvm`. (If you've just built it, run `stack clean` first to ensure it really gets rebuilt with LLVM.)
* Wait patiently
* Enjoy the result!

You don't necessarily have to use LLVM at all. However, if you can acquire binaries of the right LLVM version, that will give you some speedups.

## Usage
When `blackstar` has been built with `stack`, you can run it with
```
stack exec blackstar -- [-p|--preview] [-f|--force] [-o|--output=PATH] [-s|--starmap=PATH] SCENENAME
```
Notice the two dashes (`--`) which are required to terminate `stack`'s argument list.


`cabal` users can run `blackstar` by executing
```
cabal run -- [OPTIONS] SCENENAME
```
in the root folder of the project.

Scenes are defined using YAML config files. Look in the `scenes` folder for examples. To render the `default` scene to the directory `output`, run
```
stack exec blackstar -- scenes/default.yaml --output output
```
in the root directory of the project. The `--output` flag specifies the output directory. By default, `blackstar` searches for a starmap in the path `./stars.kdt`, but a different path can be specified using the `--starmap` flag.

The rendered files are named `scenename.png` and `scenename-bloomed.png`. The `--preview` flag can be used to render small-sized previews of the scene while adjusting the parameters. The `--force` flag will cause `blackstar` to overwrite output images without a prompt.

If a directory is given as the input scene path, `blackstar` searches non-recursively for YAML files in that directory and tries to render them. The scenes are placed in the specified output directory.

There's also a help text which can be seen by running
```
stack exec blackstar -- --help
OR
cabal run -- --help
```

Better images can be achieved by rendering larger than the target size and then scaling down (some antialiasing is achieved). This is called supersampling and is implemented in `blackstar`. It can be enabled by setting `supersampling` to true in the YAML config file — see `scenes/default-aa.yaml` for an example.

## Animation
There is a separate YAML config format for specifying animations. For example, see [default-ani.yaml](animations/default-ani.yaml).

In the first pass, the animation file must be rendered into separate config files for each frame. The `animate` executable takes care of this. First, create a directory where the frame config files will be put.
```
mkdir frames
```
Then run `animate`:
```
stack exec animate -- animations/default-ani.yaml -o frames
```
Now you should find quite a bunch of `.yaml` files in the folder `frames`.

Make another folder for the output frames:
```
mkdir frames-out
```
Now you will be able to run `blackstar` in batch mode to render the frames:
```
stack exec blackstar -- frames -o frames-out
```
This will take quite a while.

After the frames have been rendered, generate a video from the `*.png` still with your utility of preference. You can also use my script `scripts/ffmpeg-animate`, which uses `ffmpeg`. You only need to give it the prefix of the numbered frames:
```
scripts/ffmpeg-animate frames-out/default-ani
```
The output video will be rendered to `out.mkv`.

## Profiling
Thanks to `stack`, profiling is incredibly easy. Rebuild `blackstar` by running
```
stack build --profile
```
and then run it with
```
stack exec blackstar -- scenes/default.yaml -o output +RTS -p
```
The profile will be generated to `blackstar.prof`.

## TODO
As always, there's a plenty of room for improvement. For example:

* Animation: mathematically rigorous non-stationary observer
* Arbitrary textures for accretion disk (or some cool noise generator)
* Redshifting of the accretion disk
* Preview / scene planner GUI ([fltkhs](https://hackage.haskell.org/package/fltkhs))

Pull requests are welcome! If you find some cool scenes, I'd appreciate if you contributed them to this repository.



================================================
FILE: Setup.hs
================================================
import Distribution.Simple
main = defaultMain


================================================
FILE: animations/default-ani.yaml
================================================
# Animation files are quite similar to scene files. The difference is that you
# can define multiple camera instances which are then fixed to certain points
# of time (normalized to [0, 1[). The frames in between are then generated by
# interpolating between the keyframes using the interpolation method of choice.

# Scene config is similar to the scene files
scene:
    resolution: [1920, 1080]
    bloomStrength: 0.7
    starIntensity: 0.7
    starSaturation: 0.7
    diskHSV: [180, 0.1, 1.05]
    diskOpacity: 0.95
    diskInner: 1.8
    diskOuter: 13
    supersampling: true

# The number of frames in the animation
nFrames: 375
# Interpolation method. Currently, 'linear' is the only option
interpolation: 'linear'

# Specify any number of keyframes you want
keyframes:
- time: 0  # The normalize time of the keyframe. 0 is the first frame, 1 is the last
  camera:  # Camera config is similar to the scene files
    position:   [3, 3, -20]
    lookAt:     [-7, 5, 0]
    upVec:      [-0.2, 1, 0]
    fov:        1.5
- time: 1
  camera:
    position:   [-15, 1, -20]
    lookAt:     [13, -7, 0]
    upVec:      [-0.2, 1, 0]
    fov:        2


================================================
FILE: app/Animate.hs
================================================
{-# LANGUAGE DeriveDataTypeable #-}
{-# OPTIONS_GHC -fno-cse #-}

module Main where

import Prelude hiding (writeFile)
import System.Directory
import System.FilePath
import Data.Yaml (decodeFileEither, prettyPrintParseException, encode)
import System.Console.CmdArgs
import Control.Monad (forM_)
import Data.ByteString.Lazy (fromStrict, writeFile)

import Animation
import Util

data Animate = Animate { inFile :: FilePath
                       , output :: FilePath
                       , force :: Bool }
                       deriving (Show, Data, Typeable)

argparser :: Animate
argparser = Animate { inFile = def
                            &= argPos 0
                            &= typ "INPUTFILE"
                    , output = ""
                            &= help "output directory"
                            &= typ "PATH"
                    , force  = def
                            &= help "overwrite images without asking" }
                    &= summary "Animation helper for Blackstar"

main :: IO ()
main = do
    cmdline <- cmdArgs argparser

    let inPath = inFile cmdline
    let basename = takeBaseName inPath
    inputExists <- doesFileExist inPath

    outPath <- normalizePath =<< case output cmdline of
                  "" -> getCurrentDirectory
                  x  -> return x
    createDirectoryIfMissing True outPath

    if inputExists then do
        config <- decodeFileEither inPath
        case config of
            Right cfg ->
                case validateKeyframes $ keyframes cfg of
                    Right () -> do
                        let nFr = nFrames cfg
                        forM_ (zip (generateFrames cfg) [(0 :: Int), 1 ..])
                            (\(frame, idx) -> do
                                let filename = outPath </> basename ++ "_" ++
                                        padZero (nFr - 1) idx <.> ".yaml"
                                let outBl = fromStrict $ encode frame
                                if force cmdline
                                    then writeFile filename outBl
                                    else promptOverwriteFile filename 
                                           (\fname -> writeFile fname outBl)
                            )
                    Left err -> putStrLn err
            Left err -> putStrLn $ "Error when decoding config:\n" ++
                                   prettyPrintParseException err
        else putStrLn "Couldn't open input file."


================================================
FILE: app/GenerateTree.hs
================================================
module Main where

import qualified Data.ByteString.Lazy as B
import System.Environment (getArgs)

import Util
import StarMap

-- Generate and store the k-d star tree from a star catalog

main :: IO ()
main = do
    args <- getArgs
    case args of
        [infile, outfile] -> do
            outfile' <- normalizePath outfile
            infile' <- normalizePath infile
            eitherMap <- readMapFromFile infile'
            case eitherMap of
                Right stars -> do
                    putStrLn "Generating the star tree..."
                    tree <- timeAction "Building the tree"
                        $ buildStarTree stars
                    let treeBl = B.fromStrict $ treeToByteString tree
                    promptOverwriteFile outfile'
                        (\filename -> B.writeFile filename treeBl)
                    putStrLn $ "Tree saved to " ++ outfile' ++ "."
                Left  err   ->  putStrLn err
        _ -> putStrLn "Usage: generate-tree <INFILE> <OUTFILE>"


================================================
FILE: app/Main.hs
================================================
{-# LANGUAGE DeriveDataTypeable #-}
{-# OPTIONS_GHC -fno-cse #-}

module Main where

import System.Directory
import Control.Monad (when, forM_)
import Data.Yaml (decodeFileEither, prettyPrintParseException)
import System.Console.CmdArgs
import System.FilePath (takeBaseName, takeExtension, (</>), (<.>))
import Data.List (sort)
import System.Console.ANSI (clearScreen, setCursorPosition)

import Raytracer
import StarMap
import ConfigFile
import ImageFilters
import Util

data Blackstar = Blackstar { preview :: Bool
                           , output :: String
                           , force :: Bool
                           , starmap :: String
                           , inputfile :: String }
                           deriving (Show, Data, Typeable)

argparser :: Blackstar
argparser = Blackstar { preview = def
                          &= help "preview render (small size)"
                      , output = ""
                          &= help "output directory"
                          &= typ "PATH"
                      , force = def
                          &= help "overwrite images without asking"
                      , starmap = "stars.kdt"
                          &= help "path to starmap"
                          &= typ "PATH"
                      , inputfile = def
                          &= argPos 0
                          &= typ "INPUTFILE"
                      } &= summary "Blackstar v0.1"

main :: IO ()
main = do
    cmdline <- cmdArgs argparser
    etree <- readTreeFromFile $ starmap cmdline
    case etree of
        Right tree -> putStrLn "Starmap successfully read."
            >> doStart cmdline tree
        Left  err  -> putStrLn $ "Error decoding star tree: \n" ++ err

doStart :: Blackstar -> StarTree -> IO ()
doStart cmdline tree = do
    -- Resolve the output directory
    when (output cmdline /= "")
        $ createDirectoryIfMissing True (output cmdline)
    outdir <- normalizePath =<< case output cmdline of
                  "" -> getCurrentDirectory
                  x  -> return x
    createDirectoryIfMissing True outdir
    -- Resolve the input file or directory
    filename <- normalizePath $ inputfile cmdline
    isDir <- doesDirectoryExist filename
    if isDir then do
            putStrLn $ filename
                ++ " is a directory. Rendering all scenes inside it..."

            inputFiles <- map (filename </>)
                . sort . filter (\f -> takeExtension f == ".yaml")
                <$> getDirectoryContents filename

            forM_ (zip inputFiles [(1 :: Int)..]) $ \(scn, idx) -> do
                clearScreen
                setCursorPosition 0 0
                putStrLn $ "Batch mode progress: " ++ show idx ++ "/"
                    ++ show (length inputFiles)
                handleScene cmdline tree outdir scn
        else handleScene cmdline tree outdir filename

handleScene :: Blackstar -> StarTree -> String -> String -> IO ()
handleScene cmdline tree outdir filename = do
    let pvw = preview cmdline
    let sceneName = takeBaseName filename
    putStrLn $ "Reading " ++ filename ++ "..."
    cfg <- decodeFileEither filename
    let sceneName' = if pvw then "prev-" ++ sceneName else sceneName
    case cfg of
        Right config -> putStrLn "Scene successfully read."
                          >> doRender cmdline (prepareScene config pvw) tree
                               sceneName' outdir
        Left  err    -> putStrLn $ prettyPrintParseException err

prepareScene :: Config -> Bool -> Config
prepareScene cfg doPreview = let
    scn = scene cfg
    (w, h) = resolution scn
    res = 300
    newRes = if w >= h then (res, res * h `div` w) else (res * w `div` h, res)
    newScn = if doPreview then scn { resolution = newRes
                                   , supersampling = False
                                   , bloomStrength = 0 }
                          else scn
    in cfg { scene = newScn }

doRender :: Blackstar -> Config -> StarTree -> String -> String -> IO ()
doRender cmdline cfg tree sceneName outdir = do
    putStrLn $ "Rendering " ++ sceneName ++ "..."
    let scn = scene cfg
    img <- timeAction "Rendering" $ render cfg tree

    let outName = outdir </> sceneName <.> ".png"

    final <- if bloomStrength scn /= 0
        then do
            putStrLn "Applying bloom..."
            bloomed <- bloom (bloomStrength scn) (bloomDivider scn) img
            timeAction "Bloom" bloomed
        else return img

    putStrLn $ "Saving to " ++ outName ++ "..."
    if force cmdline
      then writeImg final outName
      else promptOverwriteFile outName (writeImg final)

    putStrLn "Everything done. Thank you!"


================================================
FILE: blackstar.cabal
================================================
name:                blackstar
version:             0.1.0.0
synopsis:            A black hole ray tracer
description:         A black hole ray tracer written in Haskell
homepage:            http://github.com/flannelhead/blackstar#readme
license:             BSD3
license-file:        LICENSE
author:              Sakari Kapanen
maintainer:          sakari.m.kapanen@gmail.com
copyright:           2019 Sakari Kapanen
category:            Graphics
build-type:          Simple
-- extra-source-files:
cabal-version:       >=1.10

library
  hs-source-dirs:      src
  exposed-modules:     Raytracer
                     , StarMap
                     , ConfigFile
                     , ImageFilters
                     , Util
                     , Animation
  ghc-options:         -Wall
                       -O2
  build-depends:       base >= 4.7 && < 5
                     , data-default
                     , linear
                     , lens
                     , bytestring
                     , cereal
                     , vector
                     , kdt
                     , yaml
                     , aeson
                     , directory
                     , massiv
                     , massiv-io
                     , filepath
                     , time
                     , deepseq
  default-language:    Haskell2010

executable blackstar
  hs-source-dirs:      app
  main-is:             Main.hs
  ghc-options:         -Wall -rtsopts -threaded -with-rtsopts=-N
  build-depends:       base
                     , blackstar
                     , directory
                     , yaml
                     , bytestring
                     , cmdargs
                     , filepath
                     , ansi-terminal
  default-language:    Haskell2010

executable animate
  hs-source-dirs:      app
  main-is:             Animate.hs
  ghc-options:         -Wall
  build-depends:       base
                     , blackstar
                     , directory
                     , yaml
                     , bytestring
                     , cmdargs
                     , filepath
  default-language:    Haskell2010

executable generate-tree
  hs-source-dirs:      app
  main-is:             GenerateTree.hs
  ghc-options:         -Wall
  build-depends:       base
                     , blackstar
                     , directory
                     , bytestring
  default-language:    Haskell2010

source-repository head
  type:     git
  location: https://github.com/flannelhead/blackstar


================================================
FILE: scenes/closeup.yaml
================================================
camera:
    position:   [10, 1, -2]
    lookAt:     [0, 0, 6]
    upVec:      [0, 1, 0]
    fov:        1.2

scene:
    resolution: [1280, 960]
    bloomStrength: 0.7
    starIntensity: 0.7
    starSaturation: 0.7
    diskOpacity: 0.95
    diskInner: 3
    diskOuter: 9


================================================
FILE: scenes/default-aa.yaml
================================================
camera:
    position:   [0, 1, -20]
    lookAt:     [2, 0, 0]
    upVec:      [-0.2, 1, 0]
    fov:        1.5

scene:
    resolution: [1920, 1080]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskColor: [180, 0.1, 1.05]
    diskOpacity: 0.95
    diskInner: 1.8
    diskOuter: 13
    supersampling: true


================================================
FILE: scenes/default.yaml
================================================
camera:
    # All the vectors are [x, y, z] coordinates
    position:   [0, 1, -20]  # The position of the camera
    lookAt:     [2, 0, 0]  # The point to look at
    upVec:      [-0.2, 1, 0]  # The "up" direction vector which determines the
                              # orientation of the camera
    fov:        1.5           # The tangent of the view angle

scene:
    resolution: [1920, 1080]   # [width, height] of the image
    bloomStrength: 0.15 # The strength (weight) of the bloom effect. Setting this to
                        # 0 disables it entirely
    bloomDivider: 25  # A number x such that r = image width / x is the bloom radius

    starIntensity: 0.4   # The intensity (0 = black, 1 = white) of the stars
    starSaturation: 1.5  # The color saturation of the stars

    diskOpacity: 0.95  # Opacity of the accretion disk (0 = fully transparent,
                       # 1 = fully opaque)
    diskInner: 1.8     # The inner radius of the accretion disk
    diskOuter: 13      # The outer radius of the accretion disk
    diskColor: [180, 0.1, 1.05]  # The colour of the accretion disk in the HSI color space.
                                 # H: 0..360, S: 0..1, I: 0..1

    supersampling: false  # Set this to true to enable smoothing by supersampling
                          # a 4x sized image
    stepSize: 0.3  # The size of the timestep in the simulation. Usually this value
                   # should be fine.


================================================
FILE: scenes/fartheraway.yaml
================================================
camera:
    position:   [-25, 1, -60]
    lookAt:     [-12, -4, 0]
    upVec:      [0.15, 1, 0]
    fov:        2

scene:
    resolution: [1920, 1080]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskOpacity: 0.95
    diskInner: 3
    diskOuter: 12
    supersampling: true


================================================
FILE: scenes/lensing-disk.yaml
================================================
camera:
    position:   [30, 0.4, 3]
    lookAt:     [0, 0, 0]
    upVec:      [0, 1, 0.2]
    fov:        1

scene:
    resolution: [1280, 800]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskOpacity: 0.95
    diskInner: 3
    diskOuter: 12
    supersampling: true


================================================
FILE: scenes/lensing.yaml
================================================
camera:
    position:   [30, 0.4, 3]
    lookAt:     [0, 0, 0]
    upVec:      [0, 1, 0.2]
    fov:        1

scene:
    resolution: [1600, 1200]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskOpacity: 0
    diskInner: 3
    diskOuter: 12
    supersampling: true


================================================
FILE: scenes/wideangle-disk.yaml
================================================
camera:
    position:   [-6, 1, -20]
    lookAt:     [-6, -4, 0]
    upVec:      [-0.2, 1, 0]
    fov:        3.5

scene:
    resolution: [1920, 1080]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskOpacity: 0.95
    diskInner: 2.5
    diskOuter: 12
    supersampling: true


================================================
FILE: scenes/wideangle.yaml
================================================
camera:
    position:   [20, 0, 0]
    lookAt:     [0, 0, 3.5]
    upVec:      [0, 1, 0]
    fov:        2

scene:
    resolution: [1920, 1020]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskOpacity: 0
    supersampling: true


================================================
FILE: scenes/wideangle1.yaml
================================================
camera:
    position:   [0, 0, 20]
    lookAt:     [3.5, 0, 0]
    upVec:      [0, 1, 0]
    fov:        2

scene:
    resolution: [1920, 1080]
    bloomStrength: 0.15
    starIntensity: 0.4
    starSaturation: 1.5
    diskOpacity: 0
    supersampling: true


================================================
FILE: scripts/ffmpeg-animate
================================================
#!/bin/sh

# Create YouTube quality video from stills using ffmpeg
#
# USAGE:
# ffmpeg-animate PREFIX
#
# Writes to file "out.mkv" in the current directory

ffmpeg -f image2 -i "${1}_%03d.png" \
    -c:v libx264 -preset slow -crf 18 -pix_fmt yuv420p \
    -r 25 out.mkv



================================================
FILE: src/Animation.hs
================================================
{-# LANGUAGE OverloadedStrings, ScopedTypeVariables, DeriveGeneric #-}

module Animation ( Keyframe(camera, time)
                 , Animation(scene, nFrames, interpolation, keyframes)
                 , InterpolationMethod(Linear)
                 , generateFrames, validateKeyframes ) where

import Data.List (sortBy)
import Data.Ord (comparing)
import qualified ConfigFile as CF
import Data.Aeson.Types
import Linear ((*^))
import GHC.Generics

data Keyframe = Keyframe { camera :: CF.Camera
                         , time :: Double }
                         deriving (Generic)

data InterpolationMethod = Linear

data Animation = Animation { scene :: CF.Scene
                           , nFrames :: Int
                           , interpolation :: InterpolationMethod
                           , keyframes :: [Keyframe] }
                           deriving (Generic)

instance FromJSON Keyframe

instance FromJSON InterpolationMethod where
    parseJSON str = do
        (str' :: String) <- parseJSON str
        return $ case str' of
            "linear" -> Linear
            _        -> Linear

instance FromJSON Animation

validateKeyframes :: [Keyframe] -> Either String ()
validateKeyframes []  = Left "Must have at least two keyframes"
validateKeyframes [_] = validateKeyframes []
validateKeyframes frs = if time (head frs) == 0 && time (last frs) == 1
    then Right ()
    else Left "First keyframe must have time == 0, last time == 1"

generateFrames :: Animation -> [CF.Config]
generateFrames animation = let
    stepsize = (1 :: Double) / fromIntegral (nFrames animation - 1)
    -- Take the first keyframe from the scene in the config
    -- Also sort the frames by time
    frames = sortBy (comparing time) $ keyframes animation
    points = (* stepsize) . fromIntegral <$> [0 .. nFrames animation - 1]
    in map (makeFrame animation frames) points

makeFrame :: Animation -> [Keyframe] -> Double -> CF.Config
makeFrame animation frames point = let
        scn = scene animation
        mtd = interpolation animation
    in CF.Config { CF.camera = interpolate mtd frames point
                 , CF.scene = scn }

interpolate :: InterpolationMethod -> [Keyframe] -> Double -> CF.Camera
interpolate method frames t = let
        findFrames (fr1 : fr2 : frs) = if t >= time fr1 && t < time fr2
            then (fr1, fr2)
            else findFrames (fr2 : frs)
        findFrames [fr] = (fr, fr { time = time fr + 1 } )

        (f1, f2) = findFrames frames
        t' = (t - time f1) / (time f2 - time f1)

        f :: Fractional a => (Double -> a -> a) -> a -> a -> a
        f = interpolationFunction method t'

        cam1 = camera f1
        cam2 = camera f2
    in CF.Camera { CF.fov = f (*) (CF.fov cam1) (CF.fov cam2)
                 , CF.position = f (*^) (CF.position cam1) (CF.position cam2)
                 , CF.lookAt = f (*^) (CF.lookAt cam1) (CF.lookAt cam2)
                 , CF.upVec = f (*^) (CF.upVec cam1) (CF.upVec cam2) }

interpolationFunction :: Fractional a => InterpolationMethod -> Double
                                      -> (Double -> a -> a)
                                      -> a -> a -> a
{-# INLINE interpolationFunction #-}
interpolationFunction method t times a b = case method of
    Linear -> a + t `times` (b - a)


================================================
FILE: src/ConfigFile.hs
================================================
{-# LANGUAGE OverloadedStrings, FlexibleInstances, DeriveGeneric #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module ConfigFile
    ( Scene( Scene, safeDistance, stepSize, bloomStrength, bloomDivider
           , starIntensity, starSaturation, supersampling
           , diskColor, diskOpacity, diskInner, diskOuter, resolution )
    , Camera( Camera, position, lookAt, upVec, fov )
    , Config( Config, camera, scene ) ) where

import Data.Aeson.Types
import Linear
import GHC.Generics
import Graphics.ColorSpace

data Config = Config { scene :: Scene
                     , camera :: Camera }
                     deriving (Generic)

data Scene = Scene { safeDistance :: !Double
                   , stepSize :: !Double
                   , bloomStrength :: !Double
                   , bloomDivider :: !Int
                   , starIntensity :: !Double
                   , starSaturation :: !Double
                   , diskColor :: !(Pixel HSI Double)
                   , diskOpacity :: !Double
                   , diskInner :: !Double
                   , diskOuter :: !Double
                   , resolution :: !(Int, Int)
                   , supersampling :: !Bool }
                   deriving (Generic)

data Camera = Camera { position :: !(V3 Double)
                     , lookAt :: !(V3 Double)
                     , upVec :: !(V3 Double)
                     , fov :: !Double }
                     deriving (Generic)

instance FromJSON (V3 Double) where
    parseJSON vec = do
        [x, y, z] <- parseJSON vec
        return $ V3 x y z

instance ToJSON (V3 Double) where
    toJSON (V3 x y z) = toJSON [x, y, z]

instance FromJSON (Pixel HSI Double) where
    parseJSON hsi = do
        [x, y, z] <- parseJSON hsi
        return $ PixelHSI (x / 360) y z

instance ToJSON (Pixel HSI Double) where
    toJSON (PixelHSI h s i) = toJSON [360 * h, s, i]

instance FromJSON Config

instance ToJSON Config where
    toEncoding = genericToEncoding defaultOptions

instance FromJSON Camera

instance ToJSON Camera where
    toEncoding = genericToEncoding defaultOptions

instance FromJSON Scene where
    parseJSON (Object v) = Scene 0                        <$>
                           v .:? "stepSize"       .!= 0.3 <*>
                           v .:? "bloomStrength"  .!= 0.4 <*>
                           v .:? "bloomDivider"   .!= 25  <*>
                           v .:? "starIntensity"  .!= 0.7 <*>
                           v .:? "starSaturation" .!= 0.7 <*>
                           v .:? "diskColor"
                             .!= PixelHSI 0.16 0.1 0.95   <*>
                           v .:? "diskOpacity"    .!= 0   <*>
                           v .:? "diskInner"      .!= 3   <*>
                           v .:? "diskOuter"      .!= 12  <*>
                           v .:? "resolution"     .!= (1280, 720) <*>
                           v .:? "supersampling"  .!= False

    parseJSON invalid = typeMismatch "Object" invalid

instance ToJSON Scene where
    toEncoding = genericToEncoding defaultOptions


================================================
FILE: src/ImageFilters.hs
================================================
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE Strict #-}

module ImageFilters (bloom, supersample) where

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import Data.Massiv.Array
import Data.Massiv.Array.Manifest.Vector
import Data.Massiv.Array.Unsafe
import Data.Massiv.Array.IO
import Control.Monad (replicateM_)
import Control.Applicative (liftA2)
import Graphics.ColorSpace

ix1d :: Int -> Int -> Int -> Int
{-# INLINE ix1d #-}
ix1d !w !y !x = y*w + x

add :: Pixel RGB Double -> Pixel RGB Double -> Pixel RGB Double
add = liftA2 (+)
sub :: Pixel RGB Double -> Pixel RGB Double -> Pixel RGB Double
sub = liftA2 (-)
mul :: Double -> Pixel RGB Double -> Pixel RGB Double
mul a = fmap (a *)

boxBlur :: Int -> Int -> Image U RGB Double -> IO (Image U RGB Double)
boxBlur !r !passes img = let
    myDims@(h :. w) = size img
    rows' = U.enumFromN (0 :: Int) h
    cols' = U.enumFromN (0 :: Int) w

    -- Functions to safely index a vector representing an image with specialized
    -- horizontal / vertical bound checks. Out of bounds indices return a black
    -- pixel.
    {-# INLINE ixh #-}
    {-# INLINE ixv #-}
    {-# INLINE ix1d' #-}
    ix1d' = ix1d w
    ixh v y x
        | x < 0 || x >= w = PixelRGB 0 0 0
        | otherwise = U.unsafeIndex v $ ix1d' y x
    ixv v x y
        | y < 0 || y >= h = PixelRGB 0 0 0
        | otherwise = U.unsafeIndex v $ ix1d' y x

    -- Normalize by the "width" of the kernel
    normFactor :: Double
    {-# INLINE normFactor #-}
    normFactor = 1 / (2*fromIntegral r + 1)

    {-# INLINE blur #-}
    blur writeToVec crds ix1df readf vecIn y = let
        {-# INLINE pix #-}
        -- A function to yield a pixel from the image vector
        pix = readf vecIn y
        -- Starting value
        startVal = U.foldl1' add . U.map pix . U.unsafeTake r $ crds
        {-# INLINE accumulate #-}
        accumulate !rgb x = do
            let newRGB =  (rgb `add` pix (x+r)) `sub` pix (x-r)
            _ <- writeToVec (ix1df y x) $ mul normFactor newRGB
            return newRGB
        -- Sweep over the row / col of the image
        in U.foldM'_ accumulate startVal crds
    in do
        mv <- U.thaw $ toVector img
        let wrt = MU.unsafeWrite mv
        replicateM_ passes $ do
            -- First blur horizontally
            tmp1 <- U.freeze mv
            U.mapM_ (blur wrt cols' ix1d' ixh tmp1) rows'
            -- Then vertically
            tmp2 <- U.freeze mv
            U.mapM_ (blur wrt rows' (flip ix1d') ixv tmp2) cols'
        out <- U.unsafeFreeze mv
        return $ fromVector Par myDims out

bloom :: Double -> Int -> Image U RGB Double -> IO (Image U RGB Double)
bloom strength divider img = do
    let myDims@(_ :. w) = size img
    blurred <- boxBlur (w `div` divider) 3 img
    return . makeArrayR U Par myDims
        $ \ix -> img `unsafeIndex` ix `add`
                 mul strength (blurred `unsafeIndex` ix)

supersample :: Image U RGB Double -> Image U RGB Double
supersample img = let
    h :. w = size img
    {-# INLINE pix #-}
    pix y x = img `unsafeIndex` (y :. x)
    {-# INLINE f #-}
    f (y :. x) = mul 0.25
        $ pix (2*y) (2*x) `add` pix (2*y+1) (2*x) `add` pix (2*y) (2*x+1)
                          `add` pix (2*y+1) (2*x+1)
    in makeArrayR U Par (Sz ((h `div` 2) :. (w `div` 2))) f


================================================
FILE: src/Raytracer.hs
================================================
{-# LANGUAGE StrictData #-}
{-# LANGUAGE Strict #-}

module Raytracer (render, writeImg) where

import Linear hiding (lookAt, mult, trace)
import qualified Linear as L
import Control.Applicative
import Control.Lens
import Data.Default
import Data.Massiv.Array as A
import Data.Massiv.Array.IO
import Graphics.ColorSpace
import Prelude as P

import StarMap
import ConfigFile
import ImageFilters

data Layer = Layer (Pixel RGBA Double) | Bottom (Pixel RGBA Double) | None
data PhotonState = PhotonState (V3 Double) (V3 Double)

sRGB :: Double -> Double
sRGB x = let
  a = 0.055
  in if x < 0.0031308 then 12.92 * x
    else (1 + a) * x ** (1.0 / 2.4) - a

writeImg :: Image U RGB Double -> FilePath -> IO ()
writeImg img path =
    writeArray PNG def path
      . A.map (toWord8 . fmap sRGB) $ img

blend :: Pixel RGBA Double -> Pixel RGBA Double -> Pixel RGBA Double
blend src@(PixelRGBA _ _ _ ta) = let
    comp tc bc = tc + bc * (1 - ta)
    in liftA2 comp src

-- Generate the sight rays ie. initial conditions for the integration
generateRay :: Config -> Ix2 -> PhotonState
generateRay cfg (y' :. x') = PhotonState vel pos
    where cam = camera cfg
          pos = position cam
          scn = scene cfg
          w = fromIntegral . fst $ resolution scn
          h = fromIntegral . snd $ resolution scn
          matr = L.lookAt pos (lookAt cam) (upVec cam) ^. _m33
          vel  = L.normalize . (L.transpose matr !*)
                 $ V3 (fov cam * (fromIntegral x' / w - 0.5))
                      (fov cam * (0.5 - fromIntegral y' / h) * h/w)
                      (-1)

render :: Config -> StarTree -> Image U RGB Double
render cfg startree = let
    scn = scene cfg
    cam = camera cfg
    (w, h) = resolution scn
    res@(w', h') = if supersampling scn then (2*w, 2*h) else (w, h)
    scn' = scn { safeDistance =
                   max (50^(2 :: Int)) (2 * quadrance (position cam))
               , diskInner = diskInner scn ^ (2 :: Int)
               , diskOuter = diskOuter scn ^ (2 :: Int)
               , resolution = res }
    cfg' = cfg { scene = scn' }
    diskRGB = toPixelRGB $ diskColor scn
    img = makeArrayR U Par (h' :. w') $ traceRay cfg' diskRGB startree :: Image U RGB Double
    in if supersampling scn then supersample img else img

traceRay :: Config -> Pixel RGB Double -> StarTree -> Ix2
            -> Pixel RGB Double
traceRay cfg diskRGB startree pt = let
        ray@(PhotonState vel pos) = generateRay cfg pt
        h2 = quadrance $ pos `cross` vel
        scn = scene cfg
    in dropAlpha . colorize scn diskRGB startree h2 $ ray

colorize :: Scene -> Pixel RGB Double -> StarTree -> Double -> PhotonState
            -> Pixel RGBA Double
colorize scn diskRGB startree h2 crd = let
    colorize' rgba crd' = let
        newCrd = rk4 (stepSize scn) h2 crd'
        in case findColor scn diskRGB startree crd' newCrd of
            Layer rgba' -> colorize' (blend rgba rgba') newCrd
            Bottom rgba' -> blend rgba rgba'
            None -> colorize' rgba newCrd
    in colorize' (PixelRGBA 0 0 0 0) crd

findColor :: Scene -> Pixel RGB Double -> StarTree -> PhotonState -> PhotonState
             -> Layer
{-# INLINE findColor #-}
findColor scn diskRGB startree (PhotonState vel pos@(V3 _ y _))
    (PhotonState _ newPos@(V3 _ y' _))
    | r2 < 1 = Bottom (PixelRGBA 0 0 0 1)  -- already passed the event horizon
    | r2 > safeDistance scn = Bottom . addAlpha 1.0
        $ starLookup startree (starIntensity scn) (starSaturation scn) vel
    | diskOpacity scn /= 0 && signum y' /= signum y
        && r2ave > diskInner scn && r2ave < diskOuter scn
        = Layer $ diskColor' scn diskRGB (sqrt r2ave)
    | otherwise = None
    where r2 = quadrance pos
          r2' = quadrance newPos
          r2ave = (y'*r2 - y*r2') / (y' - y)

diskColor' :: Scene -> Pixel RGB Double -> Double -> Pixel RGBA Double
{-# INLINE diskColor' #-}
diskColor' scn diskRGB r = let
    rInner = sqrt (diskInner scn)
    rOuter = sqrt (diskOuter scn)
    intensity = sin (pi * ((rOuter - r) / (rOuter - rInner)) ^ (2 :: Int))
    rgb = fmap (* intensity) diskRGB
    in addAlpha (intensity * diskOpacity scn) rgb

rk4 :: Double -> Double -> PhotonState -> PhotonState
{-# INLINE rk4 #-}
rk4 h h2 y = let
    mul :: Double -> PhotonState -> PhotonState
    {-# INLINE mul #-}
    mul a (PhotonState u v) = PhotonState (u ^* a) (v ^* a)

    add :: PhotonState -> PhotonState -> PhotonState
    {-# INLINE add #-}
    add (PhotonState x z) (PhotonState u v) = PhotonState (x ^+^ u) (z ^+^ v)

    f :: PhotonState -> PhotonState
    {-# INLINE f #-}
    f (PhotonState vel pos) =
        PhotonState (-1.5*h2 / (norm pos ^ (5 :: Int)) *^ pos) vel

    k1 = f y 
    k2 = f $ y `add` mul (h / 2) k1
    k3 = f $ y `add` mul (h / 2) k2
    k4 = f $ y `add` mul h k3
    sumK = k1 `add` mul 2 k2 `add` mul 2 k3 `add` k4
    in y `add` mul (h / 6) sumK


================================================
FILE: src/StarMap.hs
================================================
{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE Strict #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

module StarMap
    ( Star, StarTree, StoredStarTree
    , readMapFromFile, treeToByteString, readTreeFromFile
    , buildStarTree, starLookup ) where

import Control.Monad
import Control.Applicative (liftA2)
import Data.Word
import Data.Char
import Data.Foldable
import qualified Data.ByteString as B
import Data.Serialize as S
import Data.KdMap.Static
import Linear as L
import Graphics.ColorSpace

import Util

type Star = (V3 Double, (Int, Double, Double))
type StarTree = KdMap Double (V3 Double) (Int, Double, Double)
type StoredStar = (V3 Double, (Int, Char))
type StoredStarTree = KdMap Double (V3 Double) (Int, Char)

instance Serialize StoredStarTree
instance Serialize (TreeNode Double (V3 Double) (Int, Char))

-- We can't serialize functions but let's hack around it so that we can
-- serialize the KdMap anyway
instance Serialize (SquaredDistanceFn Double (V3 Double)) where
    put _ = put (0 :: Word8)
    get = skip 1 >> return (defaultSqrDist toList)

instance Serialize (PointAsListFn Double (V3 Double)) where
    put _ = put (0 :: Word8)
    get = skip 1 >> return toList

-- Parse the star list in the binary format specified at
-- http://tdc-www.harvard.edu/software/catalogs/ppm.entry.html
readMap :: Get [StoredStar]
readMap = do
    -- Skip the header
    skip 28
    nBytes <- remaining
    replicateM (nBytes `div` 28) $ do
        ra <- getFloat64be
        dec <- getFloat64be
        spectral <- getWord8
        skip 1
        mag <- getInt16be
        skip 8
        return ( raDecToCartesian ra dec
               , (fromIntegral mag, chr $ fromIntegral spectral) )

starColor' :: (Int, Char) -> (Int, Double, Double)
starColor' (mag, ch) = let (h, s) = starColor ch in (mag, h, s)

-- Some nice colour values for different spectral types
starColor :: Char -> (Double, Double)
starColor 'O' = (0.631, 0.39)
starColor 'B' = (0.628, 0.33)
starColor 'A' = (0.622, 0.21)
starColor 'F' = (0.650, 0.03)
starColor 'G' = (0.089, 0.09)
starColor 'K' = (0.094, 0.29)
starColor 'M' = (0.094, 0.56)
starColor _   = (0, 0)

raDecToCartesian :: Double -> Double -> V3 Double
raDecToCartesian ra dec = V3 (cos dec*cos ra) (cos dec*sin ra) (sin dec)

readMapFromFile :: FilePath -> IO (Either String [StoredStar])
readMapFromFile path = do
    ebs <- readSafe path
    return $ ebs >>= runGet readMap

readTreeFromFile :: FilePath -> IO (Either String StarTree)
readTreeFromFile path = do
    ebs <- readSafe path
    return $ fmap starColor' <$> (S.decode =<< ebs)

treeToByteString :: StoredStarTree -> B.ByteString
treeToByteString = S.encode

buildStarTree :: [StoredStar] -> StoredStarTree
buildStarTree = build toList

starLookup :: StarTree -> Double -> Double -> V3 Double -> Pixel RGB Double
{-# INLINE starLookup #-}
starLookup starmap intensity saturation vel = let
    -- The magnitude value tells about the intensity of the star. The
    -- brighter the star, the smaller the magnitude. These constants are
    -- used for adjusting the dynamics of the rendered celestial sphere.
    max_brightness = 950 -- the "maximum brightness" magnitude
    dynamic = 50         -- "dynamic range": magnitude change that doubles intensity
    w = 0.0005           -- width parameter of the gaussian function

    nvel = L.normalize vel
    stars = inRadius starmap (3 * w) nvel

    renderPixel (pos, (mag, hue, sat)) = let
        d2 = qd pos nvel
        a = log 2 / dynamic
        -- Conversion from the log magnitude scale to linear brightness
        -- and a Gaussian intensity function. This determines the apparent size
        -- and brightness of the star.
        val = (* intensity) . min 1
            . exp $ a * (max_brightness - fromIntegral mag) - d2 / (2 * w^(2 :: Int))
        in toPixelRGB $ PixelHSI hue (saturation * sat) val
    in fmap (min 1) . foldl' (liftA2 (+)) (PixelRGB 0 0 0) $ renderPixel <$> stars


================================================
FILE: src/Util.hs
================================================
module Util ( promptOverwriteFile, readSafe, normalizePath
            , timeAction, padZero ) where

import System.Directory
import System.IO
import qualified Data.ByteString as B
import Data.Time.Clock.POSIX (getPOSIXTime)
import System.FilePath
import Control.DeepSeq

readSafe :: FilePath -> IO (Either String B.ByteString)
readSafe path = do
    exists <- doesFileExist path
    if exists then Right <$> B.readFile path
              else return . Left $ "Error: file " ++ path
                  ++ " doesn't exist.\n"

promptOverwriteFile :: FilePath -> (FilePath -> IO ()) -> IO ()
promptOverwriteFile path doWrite = do
    doesExist <- doesFileExist path
    if doesExist then do
        putStr $ "Overwrite " ++ path ++ "? [y/N] "
        hFlush stdout
        answer <- getLine
        if answer == "y" || answer == "Y" then doWrite path
                                          else putStrLn "Nothing was written."
        else doWrite path

normalizePath :: FilePath -> IO FilePath
normalizePath path = (dropTrailingPathSeparator . normalise)
    <$> makeRelativeToCurrentDirectory path

timeAction :: NFData a => String -> a -> IO a
timeAction actionName value = do
    time1 <- (round <$> getPOSIXTime) :: IO Int
    let res = value
    time2 <- round <$> (res `deepseq` getPOSIXTime)
    let secs = time2 - time1
    putStrLn $ actionName ++ " completed in " ++ show (secs `div` 60)
        ++ " min " ++ show (secs `rem` 60) ++ " sec."
    return res

padZero :: Int -> Int -> String
padZero maxVal val = let
    nDigits x = (floor . logBase 10 $ (fromIntegral x :: Double)) + 1
    nZeros = nDigits maxVal - nDigits val
    zeros = replicate nZeros '0'
    in zeros ++ show val


================================================
FILE: stack.yaml
================================================
resolver: lts-13.16
flags: {}
extra-package-dbs: []
packages:
- '.'