diff --git a/src/robofish/evaluate/evaluate.py b/src/robofish/evaluate/evaluate.py index 6670645a03ce8a6460a6e4febbbc63b3c8e32e74..1a971331217ab2154519cf79850ecf46133f043e 100644 --- a/src/robofish/evaluate/evaluate.py +++ b/src/robofish/evaluate/evaluate.py @@ -123,6 +123,7 @@ def evaluate_turn( for k, files in enumerate(files_per_path): path_turns = [] for p, file in files.items(): + # TODO: Use new io functions (entity_poses_calc_ori_rad) poses = file.select_entity_poses( None if predicate is None else predicate[k] ) diff --git a/src/robofish/io/entity.py b/src/robofish/io/entity.py index e407989275322cb24fb20e62cef184e143166788..3759ceaf57d6565808d407ff66f93dad67d45ea5 100644 --- a/src/robofish/io/entity.py +++ b/src/robofish/io/entity.py @@ -138,17 +138,22 @@ class Entity(h5py.Group): return self["orientations"] @property - def orientations_calculated(self): + def poses_calc_ori_rad(self): + # Diff between positions [t - 1, 2] diff = np.diff(self.positions, axis=0) - angles = np.arctan2(diff[:, 1], diff[:, 0]) - return angles[:, np.newaxis] - @property - def poses_calc_ori_rad(self): - return np.concatenate( - [self.positions[:-1], self.orientations_calculated], axis=1 + # angles [t - 1] + angles = utils.limit_angle_range( + np.arctan2(diff[:, 1], diff[:, 0]), _range=(0, 2 * np.pi) + ) + + # Positions with angles. The first position is cut of, as it does not have an orientation. + poses_with_calculated_orientation = np.concatenate( + [self.positions[1:], angles[:, np.newaxis]], axis=1 ) + return poses_with_calculated_orientation + @property def poses(self): return np.concatenate([self.positions, self.orientations], axis=1) @@ -157,7 +162,9 @@ class Entity(h5py.Group): def poses_rad(self): poses = self.poses # calculate the angles from the orientation vectors, write them to the third row and delete the fourth row - ori_rad = np.arctan2(poses[:, 3], poses[:, 2]) + ori_rad = utils.limit_angle_range( + np.arctan2(poses[:, 3], poses[:, 2]), _range=(0, 2 * np.pi) + ) return np.concatenate([poses[:, :2], ori_rad[:, np.newaxis]], axis=1) @property @@ -173,10 +180,12 @@ class Entity(h5py.Group): We assume, that the entity is oriented "correctly" in the first pose. So the first turn angle is 0. """ - 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)) + # poses with calulated orientation have first position cut of as it does not have an orientation + # (t - 1, (x ,y, ori)) + poses_calc_ori = self.poses_calc_ori_rad + + # Differences cuts of last item (t - 2, (dx, dy, d ori)) + diff = np.diff(poses_calc_ori, axis=0) + speed = np.linalg.norm(diff[:, :2], axis=1) + turn = utils.limit_angle_range(diff[:, 2], _range=(-np.pi, np.pi)) return np.stack([speed, turn], axis=-1) diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py index 9ff216c42b3bda4fd28c0aaade8d400d67254475..616eb9257b9628379f78fbda46cbe5b82df14713 100644 --- a/src/robofish/io/file.py +++ b/src/robofish/io/file.py @@ -143,6 +143,8 @@ class File(h5py.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. strict_validate: optional boolean, if the file should be strictly validated, before saving. The default is True. + Returns: + The file itself, so something like f = robofish.io.File().save_as("file.hdf5") works """ self.validate(strict_validate=strict_validate) @@ -155,6 +157,8 @@ class File(h5py.File): shutil.copyfile(Path(self.filename).resolve(), path) + return self + def create_sampling( self, name: str = None, diff --git a/tests/robofish/io/test_entity.py b/tests/robofish/io/test_entity.py index 7427f5aa4848715a8650b905fdfd0e1866712d14..17fdcd259e49735c1cbea607ef2cd9ae67be5d7b 100644 --- a/tests/robofish/io/test_entity.py +++ b/tests/robofish/io/test_entity.py @@ -40,13 +40,13 @@ def test_entity_turn_speed(): ) e = f.create_entity("fish", positions=positions) speed_turn = e.speed_turn - assert speed_turn.shape == (99, 2) + assert speed_turn.shape == (98, 2) - # No turn in the first timestep, since initialization turns it the right way - assert speed_turn[0, 1] == 0 + # Turns and speeds shoud be all the same afterwards, since the fish swims with constant velocity and angular velocity. + assert (np.std(speed_turn, axis=0) < 0.0001).all() - # 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[1:], axis=0) < 0.0001).all() + # Cut off the first position as it cannot be simulated + positions = positions[1:] # Use turn_speed to generate positions gen_positions = np.zeros((positions.shape[0], 3))