Fix formatting of very long arrays in generated XML

alimuldal · alimuldal · commit 462afda1e991 · 2018-11-12T18:02:04.000Z
`np.array2string` truncates very long arrays (such as mesh vertices) by inserting '...' to represent the skipped elements. We now use `np.savetxt` instead, with a more precise string representation for floats.

PiperOrigin-RevId: 220796989
diff --git a/dm_control/mjcf/attribute.py b/dm_control/mjcf/attribute.py
@@ -23,14 +23,14 @@
 import collections
 import hashlib
 import os
-import re
 
 from dm_control.mjcf import base
 from dm_control.mjcf import constants
 from dm_control.mjcf import debugging
 from dm_control.mujoco.wrapper import util
 import numpy as np
 import six
+
 from dm_control.utils import io as resources
 
 
@@ -190,10 +190,12 @@ def to_xml_string(self, prefix_root=None):  # pylint: disable=unused-argument
     if self._value is None:
       return None
     else:
-      return re.sub(r'\s+', r' ',
-                    np.array2string(
-                        self._value, precision=15,
-                        suppress_small=False, max_line_width=999)[1:-1].strip())
+      out = six.BytesIO()
+      # 17 decimal digits is sufficient to represent a double float without loss
+      # of precision.
+      # https://en.wikipedia.org/wiki/IEEE_754#Character_representation
+      np.savetxt(out, self._value, fmt='%.17g', newline=' ')
+      return util.to_native_string(out.getvalue())[:-1]  # Strip trailing space.
 
   def _check_shape(self, array):
     actual_length = array.shape[0]
diff --git a/dm_control/mjcf/attribute_test.py b/dm_control/mjcf/attribute_test.py
@@ -191,11 +191,22 @@ def testFloatArray(self):
     # XML string should not be affected by global print options
     np.set_printoptions(precision=3, suppress=True)
     mujoco.optional.float_array = [np.pi, 2, 1e-16]
-    self.assertXMLStringEqual(
-        mujoco.optional, 'float_array',
-        '3.141592653589793e+00 2.000000000000000e+00 1.000000000000000e-16')
+    self.assertXMLStringEqual(mujoco.optional, 'float_array',
+                              '3.1415926535897931 2 9.9999999999999998e-17')
     self.assertCanBeCleared(mujoco.optional, 'float_array')
 
+  def testFormatVeryLargeArray(self):
+    mujoco = self._mujoco
+    array = np.arange(2000, dtype=np.double)
+    mujoco.optional.huge_float_array = array
+    xml_string = mujoco.optional.get_attribute_xml_string('huge_float_array')
+    self.assertNotIn('...', xml_string)
+    # Check that array <--> string conversion is a round trip.
+    mujoco.optional.huge_float_array = None
+    self.assertIsNone(mujoco.optional.huge_float_array)
+    mujoco.optional.huge_float_array = xml_string
+    np.testing.assert_array_equal(mujoco.optional.huge_float_array, array)
+
   def testIntArray(self):
     mujoco = self._mujoco
     mujoco.optional.int_array = [2, 2]
@@ -393,7 +404,7 @@ def testFileNameTrimming(self):
     asset = attribute.Asset(
         contents='', extension=extension, prefix=original_filename)
     vfs_filename = asset.get_vfs_filename()
-    self.assertEqual(len(vfs_filename), constants.MAX_VFS_FILENAME_LENGTH)
+    self.assertLen(vfs_filename, constants.MAX_VFS_FILENAME_LENGTH)
 
     vfs = types.MJVFS()
     mjlib.mj_defaultVFS(vfs)
diff --git a/dm_control/mjcf/element_test.py b/dm_control/mjcf/element_test.py
@@ -340,7 +340,10 @@ def testAttach(self):
                           parent=attachment_site.parent, root=submujoco)
     self.assertEqual(
         subsubmodel_frame.to_xml_string().split('\n')[0],
-        '<body pos="0.1 0.1 0.1" quat="0. 1. 0. 0." name="subsubmodel/">')
+        '<body '
+        'pos="0.10000000000000001 0.10000000000000001 0.10000000000000001" '
+        'quat="0 1 0 0" '
+        'name="subsubmodel/">')
     self.assertEqual(subsubmodel_frame.all_children(),
                      subsubmujoco.worldbody.all_children())
 
@@ -485,9 +488,12 @@ def testAttachmentFrames(self):
     subsubmujoco_frame = submujoco.find('attachment_frame', 'subsubmodel')
     subsubmujoco_frame_xml = subsubmujoco_frame.to_xml_string(
         pretty_print=False, prefix_root=mujoco.namescope)
-    self.assertStartsWith(subsubmujoco_frame_xml,
-                          '<body pos="0.1 0.1 0.1" quat="0. 1. 0. 0." '
-                          'name="submodel/subsubmodel/">')
+    self.assertStartsWith(
+        subsubmujoco_frame_xml,
+        '<body '
+        'pos="0.10000000000000001 0.10000000000000001 0.10000000000000001" '
+        'quat="0 1 0 0" '
+        'name="submodel/subsubmodel/">')
     self.assertEqual(subsubmujoco_frame.full_identifier,
                      'submodel/subsubmodel/')
     with self.assertRaisesRegexp(AttributeError, 'not a valid child'):
@@ -497,7 +503,7 @@ def testAttachmentFrames(self):
         pretty_print=False, prefix_root=mujoco.namescope)
     self.assertEqual(
         hinge_joint_xml,
-        '<joint class="submodel/" type="hinge" axis="1. 2. 3." '
+        '<joint class="submodel/" type="hinge" axis="1 2 3" '
         'name="submodel/subsubmodel/"/>')
     self.assertEqual(hinge_joint.full_identifier, 'submodel/subsubmodel/')
 
@@ -737,7 +743,7 @@ def testFindAll(self):
     mujoco.attach(submujoco)
 
     geoms = mujoco.find_all('geom')
-    self.assertEqual(len(geoms), 6)
+    self.assertLen(geoms, 6)
     self.assertEqual(geoms[0].root, mujoco)
     self.assertEqual(geoms[1].root, mujoco)
     self.assertEqual(geoms[2].root, submujoco)
@@ -746,11 +752,9 @@ def testFindAll(self):
     self.assertEqual(geoms[5].root, submujoco)
 
     b_0 = submujoco.find('body', 'b_0')
-    self.assertEqual(len(b_0.find_all('joint')), 6)
-    self.assertEqual(
-        len(b_0.find_all('joint', immediate_children_only=True)), 1)
-    self.assertEqual(
-        len(b_0.find_all('joint', exclude_attachments=True)), 2)
+    self.assertLen(b_0.find_all('joint'), 6)
+    self.assertLen(b_0.find_all('joint', immediate_children_only=True), 1)
+    self.assertLen(b_0.find_all('joint', exclude_attachments=True), 2)
 
   def testFindAllFrameJoints(self):
     root_model = parser.from_path(_TEST_MODEL_XML)
diff --git a/dm_control/mjcf/test_assets/attribute_test_schema.xml b/dm_control/mjcf/test_assets/attribute_test_schema.xml
@@ -9,6 +9,7 @@
         <attribute name="int" type="int"/>
         <attribute name="string" type="string"/>
         <attribute name="float_array" type="array" array_type="float" array_size="3"/>
+        <attribute name="huge_float_array" type="array" array_type="float"/>
         <attribute name="int_array" type="array" array_type="int" array_size="2"/>
         <attribute name="keyword" type="keyword" valid_values="Alpha Beta Gamma"/>
         <attribute name="reftype" type="keyword" valid_values="entity optional"/>
