From 91b21846ca02d3a2a8db79d28a856ae551886049 Mon Sep 17 00:00:00 2001
From: klowrey <kendall.lowrey@gmail.com>
Date: Thu, 23 Sep 2021 14:22:35 -0700
Subject: [PATCH] Removed AxisArrays in favor of ComponentArrays for
 significantly reduced first-run latency. Note that accessing named versions
 of jlModel and jlData may incur slight performance hit if used as naive
 getindex(q, :root); escaping the name Symbol with Val, i.e. getindex(q,
 Val(:root)) informs the compiler how to better optimize.

---
 Project.toml     |  4 ++--
 README.md        | 10 +++++++---
 src/MuJoCo.jl    |  2 +-
 src/jldata.jl    |  2 +-
 src/metadata.jl  |  2 +-
 src/namify.jl    | 24 ++++++++++++++----------
 test/runtests.jl |  1 +
 7 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/Project.toml b/Project.toml
index 81d5612..015c949 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,13 +1,13 @@
 name = "MuJoCo"
 uuid = "93189219-7048-461c-94ec-443a161ed927"
 authors = ["Colin Summers", "Kendall Lowrey"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
-AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MuJoCo_jll = "32af7c3b-80ec-5621-8194-2f6cb2280831"
diff --git a/README.md b/README.md
index 852f2aa..0947966 100644
--- a/README.md
+++ b/README.md
@@ -43,7 +43,7 @@ All functions are documented for convenience:
 help?> mj_step
 ```
 
-For more complete about the MuJoCo physics simulator, see [MuJoCo's documentation](http://www.mujoco.org/book).
+For more complete info about the MuJoCo physics simulator, see [MuJoCo's documentation](http://www.mujoco.org/book).
 
 
 ### MuJoCo
@@ -67,6 +67,10 @@ m = jlModel("humanoid.xml")
 @assert m.opt.timestep == 0.002
 @set!! m.opt.timestep = 0.001
 @assert m.opt.timestep == 0.001
+d = jlData(m)
+mn, dn = namify(m, d) # Produced NamedTuples with struct fields and ComponentArrays
+dn.qpos[:root_x]      # allows users to access MuJoCo fields with names specified in the XML.
+dn.qpos[Val(:root_x)] # For better performance, create a Val of the symbol name.
 ```
 
 #### Globals
@@ -99,8 +103,8 @@ you will need to add `Lyceum/LyceumRegistry`.
 From the Julia REPL, type `]` to enter Pkg mode:
 ```julia-repl
 julia> ]
-(v1.3) pkg> registry add https://github.com/Lyceum/LyceumRegistry.git
-(v1.3) pkg> add MuJoCo
+(v1.6) pkg> registry add https://github.com/Lyceum/LyceumRegistry.git
+(v1.6) pkg> add MuJoCo
 ```
 
 Below we simulate passive dynamics and print out joint positions
diff --git a/src/MuJoCo.jl b/src/MuJoCo.jl
index 9873f53..8a05c59 100644
--- a/src/MuJoCo.jl
+++ b/src/MuJoCo.jl
@@ -2,9 +2,9 @@ module MuJoCo
 
 using MacroTools: combinestructdef
 using Base: RefValue
-using AxisArrays: Axis, AxisArray
 using
     UnsafeArrays,
+    ComponentArrays,
     StaticArrays,
     Reexport,
     BangBang
diff --git a/src/jldata.jl b/src/jldata.jl
index dee3a3e..919c7f7 100644
--- a/src/jldata.jl
+++ b/src/jldata.jl
@@ -118,4 +118,4 @@ function build_jlData()
     end
 end
 
-@eval $(build_jlData())
\ No newline at end of file
+@eval $(build_jlData())
diff --git a/src/metadata.jl b/src/metadata.jl
index 563ce96..a1f22ca 100644
--- a/src/metadata.jl
+++ b/src/metadata.jl
@@ -431,4 +431,4 @@ const STRUCT_INFO = Dict{Symbol,Dict{Symbol,<:FieldInfo}}(
         :actuator_length => FieldInfo((:nu,)),
         :efc_R => FieldInfo((:njmax,)),
     ),
-)
\ No newline at end of file
+)
diff --git a/src/namify.jl b/src/namify.jl
index 9cadfd3..e081795 100644
--- a/src/namify.jl
+++ b/src/namify.jl
@@ -219,25 +219,28 @@ end
 
 function namify_array(A, newaxes)
     axisarray_axes = map(newaxes) do newax
-        axname = newax.name
-        if newax isa UnnamedAxis
-            Axis{axname}(1:length(newax))
+        l = length(newax)
+        if newax isa LinearAxis 
+            keys = Tuple(newax.labels)
+            vals = Tuple(1:l)
         else
-            Axis{axname}(SVector{length(newax),Symbol}(newax.labels))
+            keys = (newax.name,)
+            vals = (1:l,)
         end
+        nt = (; zip(keys, vals)...)
+        ComponentArrays.Axis(nt)
     end
-    newsize = length.(axisarray_axes)
-    AxisArray(reshape(A, newsize), axisarray_axes)
+    newsize = map(length, newaxes)
+    ComponentArray{axisarray_axes}(reshape(A, newsize))
 end
 
 function namify_one(x::Union{jlModel,jlData}, axismapping)
     newfields = []
     fieldnames = Symbol[]
-    for fieldname in propertynames(x)
+    for (fieldname, axis) in axismapping
         field = getproperty(x, fieldname)
-        if field isa Array && haskey(axismapping, fieldname)
-            newaxes = axismapping[fieldname]
-            push!(newfields, namify_array(field, newaxes))
+        if field isa Array
+            push!(newfields, namify_array(field, axis))
             push!(fieldnames, fieldname)
         end
     end
@@ -250,3 +253,4 @@ function namify(m::jlModel, d::jlData)
     m_axmappings = associate_axes(m, m, sz2ax)
     namify_one(m, m_axmappings), namify_one(d, d_axmappings)
 end
+
diff --git a/test/runtests.jl b/test/runtests.jl
index e5ab411..39251d7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -34,6 +34,7 @@ end
 
 function test_uninit(T)
     zeroarg = T()
+    zeroarg = T() # huh? first empty init may be different
     posonly = T((getfield(zeroarg, name) for name in fieldnames(T))...)
     isequal(zeroarg, posonly) && hash(zeroarg) == hash(posonly)
 end