
   0 | module Geom.Gen2D.State
   1 | 
   2 | import Data.Array.Mutable
   3 | import public Data.Graph.Indexed
   4 | import public Data.Linear.List
   5 | import public Data.Linear.Ref1
   6 | import public Data.Linear.Traverse1
   7 | import public Data.Refined
   8 | import public Geom
   9 | import public Syntax.T1
  10 | 
  11 | %default total
  12 | 
  13 | export
  14 | BOND_LEN : Double
  15 | BOND_LEN = 1.25
  16 | 
  17 | ||| Internal state for iteratively placing (parts of) the atoms
  18 | ||| in a molecule.
  19 | |||
  20 | ||| It keeps track of the atoms that have been placed so far
  21 | ||| but also of the current center of the placed atoms.
  22 | export
  23 | record PlaceST (s : Type) (k : Nat) where
  24 |   [noHints]
  25 |   constructor PST
  26 |   pos    : MArray s k (Maybe $ Point Mol)
  27 |   psum   : Ref s MolPoint
  28 |   placed : Ref s Nat
  29 | 
  30 | ||| Creates a new state for coordinate generation with all
  31 | ||| coordinates being currently unset and the current center
  32 | ||| at the origin.
  33 | export
  34 | placeST : (k : Nat) -> F1 s (PlaceST s k)
  35 | placeST k = T1.do
  36 |   m <- marray1 k Nothing
  37 |   s <- ref1 (P 0 0)
  38 |   p <- ref1 Z
  39 |   pure (PST m s p)
  40 | 
  41 | ||| Sets the coordinate of the given node.
  42 | |||
  43 | ||| This can also be used to re-place a node that was already
  44 | ||| placed. The center of all placed nodes is updated accordingly.
  45 | export
  46 | place : PlaceST s k => Fin k -> Point Mol -> F1' s
  47 | place @{st} x p =
  48 |   Core.get st.pos x >>= \case
  49 |     Just q  => T1.do
  50 |       set st.pos x (Just p)
  51 |       mod1 st.psum $ \(P x y) => P (x + p.x - q.x) (y + p.y - q.y)
  52 |     Nothing => T1.do
  53 |       set st.pos x (Just p)
  54 |       mod1 st.placed S
  55 |       mod1 st.psum $ \(P x y) => P (x + p.x) (y + p.y)
  56 | 
  57 | ||| Marks the given node as unplaced by removing its
  58 | ||| entry from the array of coordinates and adjusting the
  59 | ||| center of placed nodes accordingly.
  60 | |||
  61 | ||| In case the node is already unset, this is a no-op.
  62 | export
  63 | unplace : PlaceST s k => Fin k -> F1' s
  64 | unplace @{st} x =
  65 |   Core.get st.pos x >>= \case
  66 |     Nothing => pure ()
  67 |     Just p  => T1.do
  68 |       set st.pos x Nothing
  69 |       mod1 st.placed pred
  70 |       mod1 st.psum $ \(P x y) => P (x - p.x) (y - p.y)
  71 | 
  72 | ||| Returns the current coordinate of the given node.
  73 | |||
  74 | ||| Returns the origin in case the node is currently unplaced.
  75 | export
  76 | nodePosition : PlaceST s k => Fin k -> F1 s MolPoint
  77 | nodePosition @{st} x t =
  78 |   case Core.get st.pos x t of
  79 |     Just p  # t => p # t
  80 |     Nothing # t => origin # t
  81 | 
  82 | ||| Returns the vector connecting the two nodes by subtracit
  83 | ||| the position of `x` from the one of `y`.
  84 | export
  85 | bondVector : PlaceST s k => (x,y : Fin k) -> F1 s MolVector
  86 | bondVector x y = T1.do
  87 |   px <- nodePosition x
  88 |   py <- nodePosition y
  89 |   pure $ py - px
  90 | 
  91 | ||| True, if the given node is currently placed.
  92 | export
  93 | isPlaced : PlaceST s k => Fin k -> F1 s Bool
  94 | isPlaced @{st} x t =
  95 |   case Core.get st.pos x t of
  96 |     Just p  # t => True # t
  97 |     Nothing # t => False # t
  98 | 
  99 | ||| Returns the center of the currently placed nodes.
 100 | export
 101 | center : PlaceST s k => F1 s MolPoint
 102 | center @{st} t =
 103 |   case read1 st.placed t of
 104 |     0 # t => P 0 0 # t
 105 |     n # t =>
 106 |      let P x y # t := read1 st.psum t
 107 |          d         := cast {to = Double} n
 108 |       in P (x/d) (y/d) # t
 109 | 
 110 | ||| Computes the center of the given nodes.
 111 | |||
 112 | ||| This uses `nodePosition` internally, therefore, if one of the
 113 | ||| has not yet been placed, it will be placed at the origin.
 114 | export
 115 | centerOf : PlaceST s k => List (Fin k) -> F1 s MolPoint
 116 | centerOf vs = center2d <$> traverse1 nodePosition vs
 117 | 
 118 | ||| Returns the current node positions in an immutable array.
 119 | |||
 120 | ||| Unplaced nodes will be put at the origin.
 121 | export
 122 | getPoints : {k : _} -> PlaceST s k -> F1 s (IArray k $ Point Mol)
 123 | getPoints st t =
 124 |   let m # t := mmap (fromMaybe origin) st.pos t
 125 |    in unsafeFreeze m t
 126 | 
 127 | ||| Adjusts the position of the given node by applying the
 128 | ||| given function.
 129 | export
 130 | adjPoint : PlaceST s k => (Point Mol -> Point Mol) -> Fin k -> F1' s
 131 | adjPoint @{st} f i = nodePosition i >>= place i . f
 132 |