diff --git a/src/pypulseq/__init__.py b/src/pypulseq/__init__.py index 498b8e59..0156b1c6 100644 --- a/src/pypulseq/__init__.py +++ b/src/pypulseq/__init__.py @@ -59,6 +59,7 @@ def round_half_up(n, decimals=0): from pypulseq.opts import Opts from pypulseq.points_to_waveform import points_to_waveform from pypulseq.rotate import rotate +from pypulseq.rotate_3d import rotate_3d from pypulseq.scale_grad import scale_grad from pypulseq.split_gradient import split_gradient from pypulseq.split_gradient_at import split_gradient_at diff --git a/src/pypulseq/rotate_3d.py b/src/pypulseq/rotate_3d.py new file mode 100644 index 00000000..fd5b45c9 --- /dev/null +++ b/src/pypulseq/rotate_3d.py @@ -0,0 +1,148 @@ +from types import SimpleNamespace +from typing import List, Union + +import numpy as np +from scipy.spatial.transform import Rotation + +from pypulseq.add_gradients import add_gradients +from pypulseq.block_to_events import block_to_events +from pypulseq.make_delay import make_delay +from pypulseq.opts import Opts +from pypulseq.scale_grad import scale_grad +from pypulseq.utils.tracing import trace, trace_enabled + + +def _get_grad_abs_mag(grad: SimpleNamespace) -> float: + if grad.type == 'trap': + return abs(grad.amplitude) + return np.max(np.abs(grad.waveform)) + + +def _rotation_to_matrix(rotation: Union[np.ndarray, list, tuple, float, Rotation]) -> np.ndarray: + """Convert the supported rotation representations to a 3x3 rotation matrix.""" + if isinstance(rotation, Rotation): + return rotation.as_matrix() + + rotation = np.asarray(rotation, dtype=float) + if rotation.shape == (3, 3): + return rotation + + values = rotation.ravel() + if values.size == 4: + # Unit quaternion with the scalar component first [w, x, y, z]; scipy expects [x, y, z, w] + return Rotation.from_quat(values[[1, 2, 3, 0]]).as_matrix() + + if values.size in (1, 2): + phi = values[0] + if not -np.pi <= phi < 2 * np.pi: + raise ValueError(f'Rotation angle phi ({phi:.2f}) is invalid. Must be within [-pi, 2*pi) radians.') + rot = Rotation.from_euler('z', phi) + if values.size == 2: + theta = values[1] + if not -np.pi <= theta <= np.pi: + raise ValueError(f'Rotation angle theta ({theta:.2f}) is invalid. Must be within [-pi, pi] radians.') + rot = rot * Rotation.from_euler('y', theta) + return rot.as_matrix() + + raise ValueError( + 'The parameter `rotation` must be a 3x3 rotation matrix, a unit quaternion [w, x, y, z], ' + 'a pair of angles [phi, theta], a single angle phi or a scipy Rotation object.' + ) + + +def rotate_3d( + *args: SimpleNamespace, + rotation: Union[np.ndarray, list, tuple, float, Rotation], + system: Union[Opts, None] = None, +) -> List[SimpleNamespace]: + """ + Rotates the corresponding gradient(s) by the provided 3D rotation. Non-gradient events are not affected. + + The rotation can be given as a 3x3 rotation matrix, a unit quaternion with the scalar component first + [w, x, y, z], a pair of polar angles [phi, theta] in radians (rotation about y by theta followed by a rotation + about z by phi), a single angle phi in radians (rotation about z) or a `scipy.spatial.transform.Rotation` object. + For a single rotation about the z-axis `pypulseq.rotate()` can be faster. + + Gradients that end up on the same axis are added together with `pypulseq.add_gradients()`. Rotated gradient + components with negligible amplitude (relative to the largest input gradient) are dropped. Events that are not + modified are passed through unchanged, including their `id` attribute if present; the `id` attribute is removed + from the rotated gradients. + + See also `pypulseq.rotate.rotate()` and `pypulseq.Sequence.sequence.add_block()`. + + Parameters + ---------- + args : SimpleNamespace + Events to be rotated. Can also be a single block (as returned by `Sequence.get_block()`), in which case + the block duration is preserved as a delay event, or a single list of events, e.g. the output of a + previous `rotate_3d()` call. + rotation : numpy.ndarray, list, tuple, float or scipy.spatial.transform.Rotation + Rotation as a 3x3 matrix, a unit quaternion [w, x, y, z], polar angles [phi, theta] in radians, a single + angle phi in radians or a scipy `Rotation` object. + system : Opts, default=Opts() + System limits, used when adding gradient components on the same axis. + + Returns + ------- + rotated_grads : list[SimpleNamespace] + Non-gradient events followed by the rotated gradient(s). Can be passed directly to `Sequence.add_block()`. + + Raises + ------ + ValueError + If `rotation` has an unsupported shape or the angles are out of range. + If more than one gradient is provided for the same axis. + If a gradient has an invalid channel. + """ + if system is None: + system = Opts.default + + rot_mat = _rotation_to_matrix(rotation) + + events = block_to_events(*args) + if len(events) == 1 and isinstance(events[0], (list, tuple)): + events = events[0] # Allow nested calls, e.g. rotate_3d(rotate_3d(...), rotation=...) + + axes = ['x', 'y', 'z'] + grads_in = {} + bypass = [] + for event in events: + if event is None: + continue + if isinstance(event, (float, int)): + bypass.append(make_delay(event)) # block_duration of a block passed as input + continue + if getattr(event, 'type', None) not in ('grad', 'trap'): + bypass.append(event) + elif event.channel not in axes: + raise ValueError(f'Invalid gradient channel. Expected one of {axes}, got: {event.channel}') + elif event.channel in grads_in: + raise ValueError(f'More than one gradient on the same axis {event.channel} provided.') + else: + grads_in[event.channel] = event + + # Measure of the relevant amplitude, used to drop negligible rotated components + fthresh = 1e-6 + thresh = fthresh * max((_get_grad_abs_mag(g) for g in grads_in.values()), default=0) + + rotated_grads = [] + for j, out_axis in enumerate(axes): + components = [] + for i, in_axis in enumerate(axes): + if in_axis not in grads_in or abs(rot_mat[j, i]) < fthresh: + continue + g = scale_grad(grad=grads_in[in_axis], scale=rot_mat[j, i]) + g.channel = out_axis + components.append(g) + + if len(components) == 0: + continue + g = components[0] if len(components) == 1 else add_gradients(grads=components, system=system) + if _get_grad_abs_mag(g) >= thresh: + rotated_grads.append(g) + + if trace_enabled(): + for grad in rotated_grads: + grad.trace = trace() + + return [*bypass, *rotated_grads] diff --git a/tests/test_rotate_3d.py b/tests/test_rotate_3d.py new file mode 100644 index 00000000..4b2efdf2 --- /dev/null +++ b/tests/test_rotate_3d.py @@ -0,0 +1,237 @@ +from copy import copy +from types import SimpleNamespace + +import numpy as np +import pypulseq as pp +import pytest +from pypulseq.rotate_3d import rotate_3d +from scipy.spatial.transform import Rotation + +gx_trap = pp.make_trapezoid('x', area=1000, duration=1e-3) +gy_trap = pp.make_trapezoid('y', area=-500, duration=1e-3) +gz_trap = pp.make_trapezoid('z', area=200, duration=1e-3) +gx_ext = pp.make_extended_trapezoid('x', amplitudes=[0, 1e5, 5e4, 0], times=[0, 1e-4, 4e-4, 6e-4]) +gy_arb = pp.make_extended_trapezoid( + 'y', amplitudes=[0, 8e4, 8e4, 0], times=[0, 1e-4, 3e-4, 5e-4], convert_to_arbitrary=True +) +rf = pp.make_block_pulse(flip_angle=np.pi / 2, duration=1e-3) +adc = pp.make_adc(num_samples=100, duration=1e-3) +delay = pp.make_delay(2e-3) + + +def rot_mat(axis, angle): # standard right-handed rotation matrix + return Rotation.from_euler(axis, angle).as_matrix() + + +def by_channel(events): # dict channel -> gradient, non-gradients ignored + grads = {} + for event in events: + if getattr(event, 'type', None) not in ('grad', 'trap'): + continue + assert event.channel not in grads, f'duplicate channel {event.channel}' + grads[event.channel] = event + return grads + + +def assert_grads_equal(a, b, rel=1e-6, abs_=1e-6): + assert a.type == b.type and a.channel == b.channel + assert a.delay == pytest.approx(b.delay, abs=1e-12) + assert a.area == pytest.approx(b.area, rel=rel, abs=abs_) + if a.type == 'trap': + for f in ('amplitude', 'rise_time', 'flat_time', 'fall_time', 'flat_area'): + assert getattr(a, f) == pytest.approx(getattr(b, f), rel=rel, abs=abs_), f + else: + np.testing.assert_allclose(a.tt, b.tt, rtol=rel, atol=1e-12) + np.testing.assert_allclose(a.waveform, b.waveform, rtol=rel, atol=abs_) + for f in ('first', 'last', 'shape_dur'): + assert getattr(a, f) == pytest.approx(getattr(b, f), rel=rel, abs=abs_), f + + +def test_identity(): + amplitudes = [g.amplitude for g in (gx_trap, gy_trap, gz_trap)] + out = rotate_3d(gx_trap, gy_trap, gz_trap, rotation=np.eye(3)) + assert len(out) == 3 + assert [g.channel for g in out] == ['x', 'y', 'z'] + for rotated, original in zip(out, (gx_trap, gy_trap, gz_trap)): + assert_grads_equal(rotated, original) + # Inputs are not modified + assert [g.amplitude for g in (gx_trap, gy_trap, gz_trap)] == amplitudes + + +def test_rotation_input_forms(): + angle = 0.3 + c, s = np.cos(angle / 2), np.sin(angle / 2) + reference = by_channel(rotate_3d(gx_trap, gy_trap, rotation=rot_mat('z', angle))) + forms = [ + rot_mat('z', angle), + [c, 0, 0, s], + angle, + [angle, 0.0], + Rotation.from_euler('z', angle), + ] + for r in forms: + out = by_channel(rotate_3d(gx_trap, gy_trap, rotation=r)) + assert set(out) == set(reference) + for channel in reference: + assert_grads_equal(out[channel], reference[channel]) + + # [phi, theta] means rotation about y by theta followed by rotation about z by phi + phi, theta = 0.4, -0.7 + a = by_channel(rotate_3d(gx_trap, gy_trap, gz_trap, rotation=[phi, theta])) + b = by_channel(rotate_3d(gx_trap, gy_trap, gz_trap, rotation=rot_mat('z', phi) @ rot_mat('y', theta))) + assert set(a) == set(b) + for channel in a: + assert_grads_equal(a[channel], b[channel]) + + +@pytest.mark.parametrize('grads', [[g] for g in (gx_trap, gy_trap, gz_trap, gx_ext, gy_arb)] + [[gx_trap, gy_trap]]) +@pytest.mark.parametrize('angle', [0.0, 0.1, np.pi / 4, np.pi / 2, np.pi, -np.pi / 3, 2 * np.pi]) +@pytest.mark.parametrize('axis', ['x', 'y', 'z']) +def test_vs_rotate(axis, angle, grads): + # pypulseq's rotate() about the y-axis currently uses the opposite handedness to a right-handed + # rotation matrix (and to MATLAB's mr.rotate); compensate here until rotate() is fixed + sign = -1 if axis == 'y' else 1 + a = by_channel(pp.rotate(*[copy(g) for g in grads], axis=axis, angle=angle)) + b = by_channel(rotate_3d(*grads, rotation=rot_mat(axis, sign * angle))) + assert set(a) == set(b) + for channel in a: + assert_grads_equal(a[channel], b[channel], rel=1e-6, abs_=1e-3) + + +def test_oblique_rotation_and_inverse(): + R = Rotation.from_rotvec(0.9 * np.array([0.5, 0.3, 0.6]) / np.linalg.norm([0.5, 0.3, 0.6])).as_matrix() + out = rotate_3d(gx_trap, rotation=R) + assert len(out) == 3 + assert all(g.type == 'trap' for g in out) + for j, g in enumerate(out): + assert g.amplitude == pytest.approx(R[j, 0] * gx_trap.amplitude, rel=1e-9) + amplitudes = np.array([g.amplitude for g in out]) + assert np.linalg.norm(amplitudes) == pytest.approx(abs(gx_trap.amplitude), rel=1e-9) + + back = rotate_3d(*out, rotation=R.T) + assert len(back) == 1 + assert back[0].channel == 'x' + assert_grads_equal(back[0], gx_trap, abs_=1e-6 * abs(gx_trap.amplitude)) + + +def test_passthrough_and_id(): + gx = copy(gx_trap) + gx.id = 7 + rf_with_id = copy(rf) + rf_with_id.id = 3 + out = rotate_3d(rf_with_id, gx, adc, delay, rotation=rot_mat('z', np.pi / 2)) + assert out[0] is rf_with_id + assert out[1] is adc + assert out[2] is delay + assert out[3].type == 'trap' and out[3].channel == 'y' + assert not hasattr(out[3], 'id') + assert rf_with_id.id == 3 + assert gx.id == 7 + assert gx.amplitude == pytest.approx(gx_trap.amplitude) + + # None inputs are skipped + out = rotate_3d(None, gx_trap, None, rotation=np.eye(3)) + assert len(out) == 1 + assert_grads_equal(out[0], gx_trap) + + +def test_block_and_list_input(): + seq = pp.Sequence() + seq.add_block(rf, gx_trap, gz_trap, adc) + block = seq.get_block(1) + R = rot_mat('x', 0.5) + a = rotate_3d(block, rotation=R) + b = rotate_3d(rf, gx_trap, gz_trap, adc, rotation=R) + # The block's block_duration is preserved as one extra delay event + delays = [e for e in a if getattr(e, 'type', None) == 'delay'] + assert len(delays) == 1 + assert delays[0].delay == pytest.approx(block.block_duration) + a_events = [e for e in a if getattr(e, 'type', None) != 'delay'] + assert len(a_events) == len(b) + assert sorted(getattr(e, 'type', None) for e in a_events) == sorted(getattr(e, 'type', None) for e in b) + ca, cb = by_channel(a_events), by_channel(b) + assert set(ca) == set(cb) + for channel in ca: + assert_grads_equal(ca[channel], cb[channel]) + + # The rotated block can be passed back to add_block with the block duration preserved + seq2 = pp.Sequence() + seq2.add_block(*a) + assert seq2.block_durations[1] == pytest.approx(seq.block_durations[1]) + + # Nested call: rotate_3d(rotate_3d(...)) equals a single rotation by the combined matrix + R1, R2 = rot_mat('z', 0.4), rot_mat('x', -0.6) + nested = by_channel(rotate_3d(rotate_3d(gx_trap, gy_trap, rotation=R1), rotation=R2)) + combined = by_channel(rotate_3d(gx_trap, gy_trap, rotation=R2 @ R1)) + assert set(nested) == set(combined) + for channel in nested: + assert_grads_equal(nested[channel], combined[channel], abs_=1e-3) + + +def test_negligible_components_dropped(): + out = rotate_3d(gx_trap, rotation=rot_mat('z', np.pi / 2)) + assert len(out) == 1 + assert out[0].channel == 'y' + assert out[0].amplitude == pytest.approx(gx_trap.amplitude, rel=1e-6) + + out = rotate_3d(gz_trap, rotation=rot_mat('z', 1.0)) + assert len(out) == 1 + assert_grads_equal(out[0], gz_trap) + + +def test_mixed_gradient_types_are_added(): + out = by_channel(rotate_3d(gx_trap, gy_arb, rotation=rot_mat('z', np.pi / 4))) + assert set(out) == {'x', 'y'} + assert out['x'].type == 'grad' and out['y'].type == 'grad' + c, s = np.cos(np.pi / 4), np.sin(np.pi / 4) + assert out['x'].area == pytest.approx(c * gx_trap.area - s * gy_arb.area, rel=1e-6) + assert out['y'].area == pytest.approx(s * gx_trap.area + c * gy_arb.area, rel=1e-6) + + +def test_errors(): + with pytest.raises(ValueError, match='More than one gradient'): + rotate_3d(gx_trap, gx_trap, rotation=np.eye(3)) + with pytest.raises(ValueError, match='Invalid gradient channel'): + rotate_3d(SimpleNamespace(type='trap', channel='q'), rotation=np.eye(3)) + with pytest.raises(ValueError, match='must be a 3x3'): + rotate_3d(gx_trap, rotation=[1, 2, 3]) + with pytest.raises(ValueError, match='phi'): + rotate_3d(gx_trap, rotation=7.0) + with pytest.raises(ValueError, match='theta'): + rotate_3d(gx_trap, rotation=[0.0, 4.0]) + with pytest.raises(TypeError): + rotate_3d(gx_trap) + + +def test_system_limits_enforced(): + sys_lim = pp.Opts(max_grad=30, grad_unit='mT/m', max_slew=100, slew_unit='T/m/s') + g1 = pp.make_trapezoid('x', amplitude=0.95 * sys_lim.max_grad, flat_time=1e-3, system=sys_lim) + g2 = pp.make_trapezoid('y', amplitude=0.95 * sys_lim.max_grad, flat_time=1e-3, system=sys_lim) + with pytest.raises(ValueError): + rotate_3d(g1, g2, rotation=rot_mat('z', np.pi / 4), system=sys_lim) + + sys_relaxed = pp.Opts(max_grad=60, grad_unit='mT/m', max_slew=200, slew_unit='T/m/s') + out = rotate_3d(g1, g2, rotation=rot_mat('z', np.pi / 4), system=sys_relaxed) + # Equal-amplitude x and y gradients rotated by 45 deg cancel on x and combine to a single y trap + assert len(out) == 1 + assert out[0].type == 'trap' and out[0].channel == 'y' + assert out[0].amplitude == pytest.approx(np.sqrt(2) * 0.95 * sys_lim.max_grad, rel=1e-6) + + +def test_sequence_roundtrip(tmp_path): + seq = pp.Sequence() + seq.add_block(*rotate_3d(rf, gx_trap, gz_trap, rotation=Rotation.from_rotvec([0.3, -0.2, 0.5]).as_matrix())) + seq.add_block(*rotate_3d(gx_ext, gy_arb, rotation=[0.7, 0.2])) + ok, err = seq.check_timing() + assert ok, err + + seq.write(tmp_path / 'r.seq') + seq2 = pp.Sequence() + seq2.read(tmp_path / 'r.seq') + + for block_id in (1, 2): + a = by_channel(list(seq.get_block(block_id).__dict__.values())) + b = by_channel(list(seq2.get_block(block_id).__dict__.values())) + assert set(a) == set(b) + for channel in a: + assert a[channel].area == pytest.approx(b[channel].area, rel=1e-4)