0 | module Geom.Gen2D.Place
  1 |
  2 | import Chem
  3 | import Geom.Gen2D.State
  4 |
  5 | %default total
  6 |
  7 | --------------------------------------------------------------------------------
  8 | -- Utilities
  9 | --------------------------------------------------------------------------------
 10 |
 11 | classMap : {k : _} -> List (Fin k, Nat) -> IArray k Nat
 12 | classMap ps =
 13 |   case sortBy (comparing snd) ps of
 14 |     []        => fill k 1
 15 |     (x,n)::ps => alloc k Z $ \m => set m x 1 >> go 1 n ps m
 16 |   where
 17 |     go : Nat -> Nat -> List (Fin k, Nat) -> MArray s k Nat -> F1 s (IArray k Nat)
 18 |     go c p []            m t = unsafeFreeze m t
 19 |     go c p ((x,n) :: xs) m t =
 20 |      let c'    := if n == p then c else S c
 21 |          _ # t := set m x c' t
 22 |       in go c' n xs m t
 23 |
 24 | export
 25 | prioritize : {k : _} -> IGraph k e n -> IArray k Nat
 26 | prioritize g = go (fill k 1) k
 27 |   where
 28 |     %inline
 29 |     compRank : IArray k Nat -> Fin k -> Nat
 30 |     compRank prev x = 3 * at prev x + sum (at prev <$> neighbours g x)
 31 |
 32 |     go : IArray k Nat -> Nat -> IArray k Nat
 33 |     go prev 0     = prev
 34 |     go prev (S n) =
 35 |      let rank := Indexed.generate k (compRank prev)
 36 |          cmap := classMap $ foldrKV (\x,v,xs => (x,v)::xs) [] rank
 37 |       in if prev == cmap then prev else go cmap n
 38 |
 39 |
 40 | parameters {k       : Nat}
 41 |            {auto dg : DebugFlag}
 42 |            {0 e, n  : Type}
 43 |            {auto ce : Cast n Elem}
 44 |            {auto ch : Cast n Hybridization}
 45 |            (g : IGraph k e n)
 46 |
 47 |   isTerminalD4 : Fin k -> Bool
 48 |   isTerminalD4 x =
 49 |    let ns := neighbours g x
 50 |     in length ns == 4 && count ((> 1) . deg g) ns <= 1
 51 |
 52 |   nextBondVector : Fin k -> MolVector -> (p,c : MolPoint) -> Bool -> MolVector
 53 |   nextBondVector x v p c trans =
 54 |     case cast @{ch} (lab g x) of
 55 |       SP => v
 56 |       _  =>
 57 |        let a  := compAngle
 58 |            va := rotate a (negate v)
 59 |            vb := rotate (negate a) (negate v)
 60 |         in if distance (translate va p) c >= distance (translate vb p) c
 61 |               then va
 62 |               else vb
 63 |
 64 |     where
 65 |       compAngle : Angle
 66 |       compAngle =
 67 |         if      isTerminalD4 x then fromDegree 45
 68 |         else if isMetal (cast $ lab g x) then fullSteps (deg g x)
 69 |         else if trans then fromDegree 120 else fromDegree 60
 70 |
 71 |   ||| Places the atoms in a linear chain.
 72 |   |||
 73 |   ||| Expects the first atom to be placed and
 74 |   ||| places the next atom according to initialBondVector. The rest of the chain
 75 |   ||| is placed such that it is as linear as possible (in the overall result, the
 76 |   ||| angles in the chain are set to 120 Deg.)
 77 |   ||| TODO: Double bond configuration
 78 |   export
 79 |   placeChain : PlaceST s k => Fin k -> List (Fin k) -> MolVector -> F1' s
 80 |   placeChain _ []      _ t = () # t
 81 |   placeChain p (n::ns) v t =
 82 |    let pp # t := nodePosition p t
 83 |        pn     := translate v pp
 84 |        b  # t := isPlaced n t
 85 |        _  # t := when1 (not b) (place n pn) t
 86 |        c  # t := State.center {k} t
 87 |     in placeChain n ns (nextBondVector n v pn c True) t
 88 |
 89 |   export
 90 |   polygonCorners :
 91 |        {auto st : PlaceST s k}
 92 |     -> List (Fin k)
 93 |     -> MolPoint
 94 |     -> (cur,step : Angle)
 95 |     -> (dir : MolVector)
 96 |     -> F1' s
 97 |   polygonCorners []        _ _   _    _   t = () # t
 98 |   polygonCorners (x :: xs) p cur step dir t =
 99 |    let theta := cur + step
100 |        p2    := translate (rotate theta dir) p
101 |        _ # t := debugIf1 "placing \{show x} at \{show p2}" t
102 |        _ # t := place x p2 t
103 |     in polygonCorners xs p theta step dir t
104 |
105 |   export
106 |   distributeAtoms : PlaceST s k => Fin k -> (us,ps : List (Fin k)) -> F1' s
107 |   distributeAtoms x []        _   t = () # t
108 |   distributeAtoms x [u]       [p] t =
109 |    let px # t := nodePosition x t
110 |        pp # t := nodePosition p t
111 |        c  # t := State.center {k} t
112 |     in placeChain x [u] (nextBondVector x (px - pp) px c True) t
113 |   distributeAtoms x us@(_::r) ps  t =
114 |    let px # t       := nodePosition x t
115 |        ps # t       := traverse1 nodePosition ps t
116 |        (start,step) := circularFreeSweep (S $ length r) px ps
117 |     in polygonCorners us px start step (V BOND_LEN 0) t
118 |
119 |   export
120 |   placeNeighbours : PlaceST s k => Fin k -> F1 s (List $ Fin k)
121 |   placeNeighbours x t =
122 |    let (us,ps) # t := partition1 isPlaced (neighbours g x) t
123 |        _       # t := debugIf1 "placing neighbours for \{show x}" t
124 |        _       # t := distributeAtoms x us ps t
125 |     in us # t
126 |
127 |   ||| Convenience method to place a single atom. This function will first find
128 |   ||| a placed neighbour does not need to be set) and then place this
129 |   ||| new atoms considering the neighbour's neighbours. Essentially this
130 |   ||| utility is useful for sprouting a new atom to an already placed
131 |   ||| structure.
132 |   export
133 |   placeAtom : PlaceST s k => Fin k -> F1' s
134 |   placeAtom x t =
135 |    let False # t := isPlaced x t | _ # t => () # t
136 |     in case filter1 isPlaced (neighbours g x) t of
137 |          []  # t => place x (P 0 0) t
138 |          [y] # t =>
139 |           let ps # t := filter1 isPlaced (neighbours g y) t
140 |            in distributeAtoms y [x] ps t
141 |          ys  # t => let c # t := centerOf ys t in place x c t
142 |