Write all of the algorithm code

This commit is contained in:
Kiana Sheibani 2022-12-11 17:51:01 -05:00
parent 73512ccd77
commit 947ae0c63a
Signed by: toki
GPG key ID: 6CB106C25E86A9F7

152
Main.idr
View file

@ -1,18 +1,156 @@
module Main module Main
import Data.List
import Data.Vect
import Data.Fuel
import Data.NumIdr import Data.NumIdr
%default total
||| A type for angles in radians.
public export public export
Angle : Type
Angle = Double
data Joint : Nat -> Type where 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 : IO ()
main = putStrLn "Hello World!" main = putStrLn "Hello World!"