def assertGraphTuplesClose(self, a, b, tol=1e-6):
a_mask = (a.edge_idx < a.nodes.shape[0]).reshape(a.edge_idx.shape + (1,))
b_mask = (b.edge_idx < b.nodes.shape[0]).reshape(b.edge_idx.shape + (1,))
a = dataclasses.replace(a, edges=a.edges * a_mask)
b = dataclasses.replace(b, edges=b.edges * b_mask)
a = dataclasses.asdict(a)
b = dataclasses.asdict(b)
self.assertAllClose(a, b)
def get_array(name, field):
assert name in geometry
geom = dataclasses.asdict(geometry[name])
assert field in geom
array = geom[field]
return _to_json({'array': _encode(array)})
def get_array_chunk(name, field, offset, size):
assert name in geometry
geom = dataclasses.asdict(geometry[name])
assert field in geom
array = geom[field]
return _to_json(
{'array_chunk': _encode(array[offset:(offset + size)])})
def get_dynamic_geometry_metadata(name):
assert name in geometry
geom = geometry[name]
geom_dict = dataclasses.asdict(geom)
geom_metadata = {
'shape': str(geom),
'fields': {},
}
for field in geom_dict:
if not isinstance(geom_dict[field], onp.ndarray):
geom_metadata[field] = geom_dict[field]
continue
if len(geom_dict[field].shape) == TYPE_DIMENSIONS[field] + 1:
geom_metadata['fields'][field] = 'dynamic'
elif len(geom_dict[field].shape) == TYPE_DIMENSIONS[field]:
geom_metadata['fields'][field] = 'static'
elif len(geom_dict[field].shape) == TYPE_DIMENSIONS[field] - 1:
geom_metadata['fields'][field] = 'global'
return _to_json(geom_metadata)