diff --git a/src/robofish/io/entity.py b/src/robofish/io/entity.py
index 669ac1746ab15dd5c24e6eae1b6986bf97c6a49b..5e8898abb893fd1f799bf1049920442fa14b663e 100644
--- a/src/robofish/io/entity.py
+++ b/src/robofish/io/entity.py
@@ -148,3 +148,25 @@ class Entity(h5py.Group):
poses[:, 2] = np.arctan2(poses[:, 3], poses[:, 2])
poses = poses[:, :3]
return poses
+
+ @property
+ def speed_turn_angle(self):
+ """Get the speed, turn and angles from the positions.
+
+ The vectors pointing from each position to the next are computed.
+ The output of the function describe these vectors.
+ Returns:
+ An array with shape (number_of_positions -1, 3).
+ The first column is the length of the vectors.
+ The second column is the turning angle, required to get from one vector to the next.
+ We assume, that the entity is oriented "correctly" in the first pose. So the first turn angle is 0.
+ The third column is the orientation of each vector.
+ """
+
+ diff = np.diff(self.positions, axis=0)
+ speed = np.linalg.norm(diff, axis=1)
+ angles = np.arctan2(diff[:, 1], diff[:, 0])
+ turn = np.zeros_like(angles)
+ turn[0] = 0
+ turn[1:] = utils.limit_angle_range(np.diff(angles))
+ return np.stack([speed, turn, angles], axis=-1)
\ No newline at end of file
diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py
index 05db9fa6dbd5ea6689dcf58d73aebdedef6e8d72..55ac2b2fdef1d36e7193402678251879586340c9 100644
--- a/src/robofish/io/file.py
+++ b/src/robofish/io/file.py
@@ -15,6 +15,7 @@
# -----------------------------------------------------------
import robofish.io
+from robofish.io.entity import Entity
import h5py
import numpy as np
@@ -27,6 +28,7 @@ import datetime
import tempfile
import uuid
import deprecation
+import types
default_format_version = np.array([1, 0], dtype=np.int32)
@@ -36,7 +38,7 @@ default_format_url = (
class File(h5py.File):
- """ Represents a RoboFish Track Format file, which should be used to store tracking data of individual animals or swarms.
+ """Represents a RoboFish Track Format file, which should be used to store tracking data of individual animals or swarms.
Files can be opened (with optional creation), modified inplace, and have copies of them saved.
"""
@@ -131,12 +133,12 @@ class File(h5py.File):
def __exit__(self, type, value, traceback):
# Check if the context was left under normal circumstances
if (type, value, traceback) == (None, None, None):
- if self.mode != "r": # No need to validate read only files (performance).
+ if self.mode != "r": # No need to validate read only files (performance).
self.validate()
self.close()
def save_as(self, path: Union[str, Path], strict_validate: bool = True):
- """ Save a copy of the file
+ """Save a copy of the file
Args:
path: path to a io file as a string or path object. If no path is specified, the last known path (from loading or saving) is used.
@@ -202,6 +204,9 @@ class File(h5py.File):
calendar_time_points = [
format_calendar_time_point(p) for p in calendar_time_points
]
+
+ for c in calendar_time_points:
+ print(type(c))
sampling.create_dataset(
"calendar_time_points",
data=calendar_time_points,
@@ -224,9 +229,33 @@ class File(h5py.File):
@property
def frequency(self):
- # NOTE: Only works if default sampling availabe and specified with frequency_hz.
- default_sampling = self["samplings"].attrs["default"]
- return self["samplings"][default_sampling].attrs["frequency_hz"]
+ common_sampling = self.common_sampling()
+ assert common_sampling is not None, "The sampling differs between entities."
+ assert (
+ "frequency_hz" in common_sampling.attrs
+ ), "The common sampling has no frequency_hz"
+ return common_sampling.attrs["frequency_hz"]
+
+ def common_sampling(
+ self, entities: Iterable["robofish.io.Entity"] = None
+ ) -> h5py.Group:
+ """Check if all entities have the same sampling.
+
+ Args:
+ entities: optional array of entities. If None is given, all entities are checked.
+ Returns:
+ The h5py group of the common sampling. If there is no common sampling, None will be returned.
+ """
+ custom_sampling = None
+ for entity in self.entities:
+ if "sampling" in entity["positions"].attrs:
+ this_sampling = entity["positions"].attrs["sampling"]
+ if custom_sampling is None:
+ custom_sampling = this_sampling
+ elif custom_sampling != this_sampling:
+ return None
+ sampling = self.default_sampling if custom_sampling is None else custom_sampling
+ return self["samplings"][sampling]
def create_entity(
self,
@@ -238,7 +267,7 @@ class File(h5py.File):
outlines: Iterable = None,
sampling: str = None,
) -> str:
- """ Creates a new single entity.
+ """Creates a new single entity.
Args:
TODO
@@ -277,7 +306,7 @@ class File(h5py.File):
outlines=None,
sampling=None,
) -> Iterable:
- """ Creates multiple entities.
+ """Creates multiple entities.
Args:
category: The common category for the entities. The canonical values are ['fish', 'robot', 'obstacle'].
@@ -314,7 +343,7 @@ class File(h5py.File):
@property
def entity_names(self) -> Iterable[str]:
- """ Getter for the names of all entities
+ """Getter for the names of all entities
Returns:
Array of all names.
@@ -330,52 +359,63 @@ class File(h5py.File):
@property
def entity_poses(self):
- return self.select_entity_poses(None)
+ return self.select_entity_property(None)
@property
def entity_poses_rad(self):
- return self.select_entity_poses(None, rad=True)
+ return self.select_entity_property(None, entity_property=Entity.poses_rad)
+
+ @property
+ def speeds_turns_angles(self):
+ return self.select_entity_property(
+ None, entity_property=Entity.speed_turn_angle
+ )
- def select_entity_poses(self, predicate=None, rad=False) -> Iterable:
- """ TODO: Rework
- Select an array of the poses of entities
+ def select_entity_poses(self, *args, ori_rad=False, **kwargs):
+ entity_property = Entity.poses_rad if ori_rad else Entity.poses
+ return self.select_entity_property(
+ *args, entity_property=entity_property, **kwargs
+ )
- If no name or category is specified, all entities will be selected.
+ def select_entity_property(
+ self,
+ predicate: types.LambdaType = None,
+ entity_property: property = Entity.poses,
+ ) -> Iterable:
+ """Get a property of selected entities.
+
+ Entities can be selected, using a lambda function.
+ The property of the entities can be selected.
Args:
- names: optional array of the names of selected entities
- category: optional selected category
+ predicate: a lambda function, selecting entities
+ (example: lambda e: e.category == "fish")
+ entity_property: a property of the Entity class (example: Entity.poses_rad)
Returns:
- An three dimensional array of all poses with the shape (entity, time, 4)
+ An three dimensional array of all properties of all entities with the shape (entity, time, property_length).
+ If an entity has a shorter length of the property, the output will be filled with nans.
"""
entities = self.entities
if predicate is not None:
entities = [e for e in entities if predicate(e)]
- max_timesteps = max([0] + [e.positions.shape[0] for e in entities])
+ assert self.common_sampling(entities) is not None
# Initialize poses output array
- pose_len = 3 if rad else 4
- poses_output = np.empty((len(entities), max_timesteps, pose_len))
- poses_output[:] = np.nan
+ properties = [entity_property.__get__(entity) for entity in entities]
- # Fill poses output array
- i = 0
- custom_sampling = None
- for entity in entities:
+ max_timesteps = max([0] + [p.shape[0] for p in properties])
- if "sampling" in entity["positions"].attrs:
- if custom_sampling is None:
- custom_sampling = entity["positions"].attrs["sampling"]
- elif custom_sampling != entity["positions"].attrs["sampling"]:
- raise Exception(
- "Multiple samplings found, preventing return of a single array."
- )
- poses = entity.poses_rad if rad else entity.poses
- poses_output[i][: poses.shape[0]] = poses
- i += 1
- return poses_output
+ property_array = np.empty(
+ (len(entities), max_timesteps, properties[0].shape[1])
+ )
+ property_array[:] = np.nan
+
+ # Fill output array
+ for i, entity in enumerate(entities):
+ property_array[i][: properties[i].shape[0]] = properties[i]
+ return property_array
@deprecation.deprecated(
deprecated_in="1.1.2",
@@ -394,6 +434,24 @@ class File(h5py.File):
predicate = None
return self.select_entity_poses(predicate)
+ def entity_turn_speed(self, predicate=None):
+ """Get an array of turns and speeds of the entities."""
+
+ entities = self.entities
+ if predicate is not None:
+ entities = [e for e in entities if predicate(e)]
+
+ assert self.common_sampling(entities) is not None
+
+ # Initialize poses output array
+ max_timesteps = max([0] + [e.positions.shape[0] for e in entities])
+ turn_speed = np.empty((len(entities), max_timesteps, 3))
+ turn_speed[:] = np.nan
+
+ for i, entity in enumerate(entities):
+ poses = entity.poses_rad if rad else entity.poses
+ poses_output[i][: poses.shape[0]] = poses
+
def validate(self, strict_validate: bool = True) -> (bool, str):
"""Validate the file to the specification.
@@ -414,7 +472,7 @@ class File(h5py.File):
return robofish.io.validate(self, strict_validate)
def to_string(self, output_format: str = "shape") -> str:
- """ The file is formatted to a human readable format.
+ """The file is formatted to a human readable format.
Args:
output_format: ['shape', 'full'] show the shape, or the full content of datasets
Returns:
@@ -424,7 +482,7 @@ class File(h5py.File):
def recursive_stringify(
obj: h5py.Group, output_format: str, level: int = 0
) -> str:
- """ This function crawls recursively into hdf5 groups.
+ """This function crawls recursively into hdf5 groups.
Datasets and attributes are directly attached, for groups, the function is recursively called again.
Args:
obj: a h5py group
diff --git a/src/robofish/io/utils.py b/src/robofish/io/utils.py
index 71396419de0211166baf6ca2e7f752666ba14d84..3be2c62eee39b1fe0e14dbbff4723f569033b0d1 100644
--- a/src/robofish/io/utils.py
+++ b/src/robofish/io/utils.py
@@ -1,5 +1,6 @@
import robofish.io
import numpy as np
+from typing import Union, Iterable
from pathlib import Path
import os
@@ -13,3 +14,25 @@ def np_array(*arrays):
def full_path(current_file, path):
return (Path(current_file).parent / path).resolve()
+
+
+def limit_angle_range(angle: Union[float, Iterable], _range=(-np.pi, np.pi)):
+ """Limit the range of an angle or array of angles between min and max
+
+ Any given angle in rad will be moved to the given range.
+ e.g. with min = -pi and max pi, 4.5*pi will be moved to be 0.5*pi.
+ Args:
+ angle: An angle or an array of angles (1D)
+ _range: optional tuple of range (min, max)
+ Returns:
+ limited angle(s) in the same form as the input (float or ndarray)
+ """
+ assert np.isclose(_range[1] - _range[0], 2 * np.pi)
+
+ def limit_one(value):
+ return (value - _range[0]) % (2 * np.pi) + _range[0]
+
+ if isinstance(angle, Iterable):
+ return np.array([limit_one(v) for v in angle])
+ else:
+ return limit_one(angle)
\ No newline at end of file
diff --git a/tests/robofish/io/test_entity.py b/tests/robofish/io/test_entity.py
index 8c4ab6946a68a316bea62d2ac90a3be56176055a..4cf14f4bc2eb204dbc03795f2a842c1ad4f71082 100644
--- a/tests/robofish/io/test_entity.py
+++ b/tests/robofish/io/test_entity.py
@@ -29,3 +29,38 @@ def test_entity_object():
assert np.isclose(poses_rad[i, 2], poses_rad_retrieved[i, 2]) or np.isclose(
poses_rad[i, 2] - 2 * np.pi, poses_rad_retrieved[i, 2]
)
+
+
+def test_entity_turn_speed():
+ f = robofish.io.File(world_size_cm=[100, 100], frequency_hz=25)
+ circle_rad = np.linspace(0, 2 * np.pi, num=100)
+ circle_size = 40
+ positions = np.stack(
+ [np.cos(circle_rad) * circle_size, np.sin(circle_rad) * circle_size], axis=-1
+ )
+ e = f.create_entity("fish", positions=positions)
+ speed_turn_angle = e.speed_turn_angle
+ assert speed_turn_angle.shape == (99, 3)
+
+ # No turn in the first timestep, since initialization turns it the right way
+ assert speed_turn_angle[0, 1] == 0
+
+ # Turns and speeds shoud afterwards be all the same afterwards, since the fish swims with constant velocity and angular velocity.
+ assert (np.std(speed_turn_angle[1:, :2], axis=0) < 0.0001).all()
+
+ # Use turn_speed to generate positions
+ gen_positions = np.zeros((positions.shape[0], 3))
+ gen_positions[0, :2] = positions[0]
+ gen_positions[0, 2] = speed_turn_angle[0, 2]
+
+ for i, (speed, turn, angle) in enumerate(speed_turn_angle):
+ new_angle = gen_positions[i, 2] + turn
+ gen_positions[i + 1] = [
+ gen_positions[i, 0] + np.cos(new_angle) * speed,
+ gen_positions[i, 1] + np.sin(new_angle) * speed,
+ new_angle,
+ ]
+
+ # The resulting positions should almost be equal to the the given positions
+ print(gen_positions[:, :2] - positions)
+ assert np.isclose(positions, gen_positions[:, :2], atol=1.0e-5).all()
diff --git a/tests/robofish/io/test_file.py b/tests/robofish/io/test_file.py
index 013d3a51056df5378bb133e2b2f5155c263b2e28..99ccbf4cdc827d1f2bccf12f0f87f7a7b5e244bc 100644
--- a/tests/robofish/io/test_file.py
+++ b/tests/robofish/io/test_file.py
@@ -89,11 +89,11 @@ def test_multiple_entities():
assert (returned_poses == poses).all()
# Just get the array for some names
- returned_poses = sf.select_entity_poses(lambda e: e.name in ["fish_1", "fish_2"])
+ returned_poses = sf.select_entity_property(lambda e: e.name in ["fish_1", "fish_2"])
assert (returned_poses == poses[:2]).all()
# Filter on both category and name
- returned_poses = sf.select_entity_poses(
+ returned_poses = sf.select_entity_property(
lambda e: e.category == "fish" and e.name == "fish_1"
)
assert (returned_poses == poses[:1]).all()
@@ -103,7 +103,7 @@ def test_multiple_entities():
obs_poses = np.random.random((agents, 1, 3))
returned_names = sf.create_multiple_entities("obstacle", poses=obs_poses)
# Obstacles should not be returned when only fish are selected
- returned_poses = sf.select_entity_poses(lambda e: e.category == "fish")
+ returned_poses = sf.select_entity_property(lambda e: e.category == "fish")
assert (returned_poses == poses).all()
# for each of the entities
@@ -139,6 +139,13 @@ def test_multiple_entities():
return sf
+def test_speeds_turns_angles():
+ with robofish.io.File(world_size_cm=[100, 100], frequency_hz=25) as f:
+ poses = np.zeros((10, 100, 3))
+ f.create_multiple_entities("fish", poses=poses)
+ assert (f.speeds_turns_angles == 0).all()
+
+
def test_broken_sampling(caplog):
sf = robofish.io.File(world_size_cm=[10, 10])
caplog.set_level(logging.ERROR)