diff --git a/Main.idr b/Main.idr index 0a2b76b..8984ee2 100644 --- a/Main.idr +++ b/Main.idr @@ -1,18 +1,156 @@ module Main +import Data.List +import Data.Vect +import Data.Fuel import Data.NumIdr +%default total + -||| A type for angles in radians. public export -Angle : Type -Angle = Double - data Joint : Nat -> Type where - Revolute : Angle -> Angle -> Joint n + Revolute : (a, b : Double) -> Joint (2 + n) + Prismatic : (a, b : Double) -> Joint (1 + n) + + +jointAction : {n : _} -> Joint n -> Double -> Maybe (Rigid n Double) +jointAction {n=S n} (Prismatic a b) x = + guard (a < x && x < b) + $> cast (translate $ vector $ x :: replicate n 0) +jointAction {n=S (S n)} (Revolute a b) x = + guard (a < x && x < b) + $> unsafeMkTrans (indexSetRange [EndBound 2,EndBound 2] + (rewrite rangeLenZ 2 in rotate2D x) identity) + + + +-- Links are directly represented by rigid transformations, i.e. rotations +-- composed with translations. The rotation component allows the link to modify +-- the orientation of the next joint, but reflections are disallowed to enforce +-- the right-hand rule for all joints. +Link : Nat -> Type +Link n = Rigid n Double + + +ArmElement : Nat -> Type +ArmElement n = List (Either (Link n) (Joint n)) + + +link : Vector n Double -> ArmElement n +link v = [Left $ cast (translate v)] + +linkX : {n : _} -> Double -> ArmElement (1 + n) +linkX x = link $ vector (x :: replicate n 0) + + +revolute2D : (a, b : Double) -> ArmElement 2 +revolute2D a b = [Right $ Revolute a b] + +revoluteX : (a, b : Double) -> ArmElement 3 +revoluteX a b = [Left $ cast (Rotation.rotate3DY (pi/2)), + Right $ Revolute a b, + Left $ cast (Rotation.rotate3DY (-pi/2))] + +revoluteY : (a, b : Double) -> ArmElement 3 +revoluteY a b = [Left $ cast (Rotation.rotate3DX (-pi/2)), + Right $ Revolute a b, + Left $ cast (Rotation.rotate3DX (pi/2))] + +revoluteZ : (a, b : Double) -> ArmElement 3 +revoluteZ a b = [Right $ Revolute a b] + +countJoints : ArmElement n -> Nat +countJoints [] = 0 +countJoints (Left _ :: xs) = countJoints xs +countJoints (Right _ :: xs) = S $ countJoints xs + +getLimits : (arm : ArmElement n) -> Vect (countJoints arm) (Double, Double) +getLimits [] = [] +getLimits (Left _ :: xs) = getLimits xs +getLimits (Right (Revolute a b) :: xs) = (a,b) :: getLimits xs +getLimits (Right (Prismatic a b) :: xs) = (a,b) :: getLimits xs + +ArmConfig : ArmElement n -> Type +ArmConfig arm = Vector (countJoints arm) Double + + +forwardTransform : {n : _} -> (arm : ArmElement n) -> ArmConfig arm + -> Maybe (Rigid n Double) +forwardTransform arm = go arm . toVect + where + go : (arm : ArmElement n) -> Vect (countJoints arm) Double + -> Maybe (Rigid n Double) + go [] _ = Just identity + go (Left l :: xs) cs = map (l *.) (go xs cs) + go (Right j :: xs) (c :: cs) = [| jointAction j c *. go xs cs |] + + +forward : {n : _} -> (arm : ArmElement n) -> ArmConfig arm -> + Maybe (Point n Double) +forward arm cs = map (fromVector . getTranslationVector . getHMatrix) + $ forwardTransform arm cs + + +Simplex : ArmElement n -> Type +Simplex arm = Vect (S $ countJoints arm) (ArmConfig arm) + +initialSimplex : (arm : ArmElement n) -> Maybe (Simplex arm) +initialSimplex arm = + let limits = getLimits arm + orig = vector $ map (\(a,b) => (a+b)/2) limits + variance = vector $ map (\(a,b) => a*0.4+b*0.6) limits + in guard (all (uncurry (<)) limits) + $> map (\case + FZ => orig + FS i => indexSet [i] (variance !! i) orig) + Fin.range + +inverse : {n : _} -> (fuel : Fuel) -> (arm : ArmElement n) -> Point n Double + -> {auto ok : IsSucc (countJoints arm)} -> Maybe (ArmConfig arm) +inverse fuel arm p = go fuel !(initialSimplex arm) + where + sndLast : {n : _} -> {auto ok : IsSucc n} -> Vect (S n) a -> a + sndLast {n=S n,ok=ItIsSucc} v = last $ init v + + cost : ArmConfig arm -> Maybe Double + cost c = map (\p' => normSq $ p -. p') (forward arm c) + + -- The standard library currently doesn't have sorting for Vects, + -- so we have to improvise a bit. + sort : Simplex arm -> Simplex arm + sort s = believe_me $ Vect.fromList $ sortBy (compare `on` cost) $ toList s + + go : Fuel -> Simplex arm -> Maybe (ArmConfig arm) + go Dry _ = Nothing + go (More fuel) simplex = do + guard (all (and . zipWith (\(a,b),x => a <= x && x <= b) + (getLimits arm) . toVect) simplex) *> + let simplex = unsafePerformIO (let s = sort simplex in printLn s $> s) + best = head simplex + cbest = !(cost best) + in if cbest < 0.00001 then Just best + else let + worst = last simplex + cworst = !(cost worst) + centroid = sum (init simplex) *. (recip $ cast {to=Double} $ countJoints arm) + costcen = !(cost centroid) + vertexr = centroid *. 2.0 - worst + cvr = !(cost vertexr) + in if cvr >= cbest && cvr < !(cost $ sndLast simplex) + then go fuel $ replaceAt last vertexr simplex + else if cvr < cbest + then let vertexe = centroid + 2.0 *. (vertexr - centroid) + in go fuel $ replaceAt last (if !(cost vertexe) < cvr then vertexe else vertexr) simplex + else let + comp = cvr < cworst + vertexc = centroid + 0.5 *. ((if comp then vertexr else worst) - centroid) + in if costcen < min cvr cworst + then go fuel $ replaceAt last vertexc simplex + else go fuel $ best :: map (\vrt => best + 0.5 *. (vrt - best)) (tail simplex) + -Robot : (n : Nat) -> Type -Robot n = List (Either (Joint n) (Vector n Double)) main : IO () main = putStrLn "Hello World!" +