Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enhance __setitem__ and __getitem__ methods to support Ellipsis indexing in CasadiLike class #120

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Enhance __setitem__ and __getitem__ methods to support Ellipsis indexing in CasadiLike class #120

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
03167f8
Enhance __setitem__ and __getitem__ methods to support Ellipsis index…
GiulioRomualdi Jan 9, 2025
3ec8e0b
Remove unnecessary blank line in CasadiLike class
GiulioRomualdi Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions src/adam/casadi/casadi_like.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,11 +85,26 @@ def __truediv__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike"

def __setitem__(self, idx, value: Union["CasadiLike", npt.ArrayLike]):
"""Overrides set item operator"""
self.array[idx] = value.array if type(self) is type(value) else value
if idx is Ellipsis:
self.array = value.array if isinstance(value, CasadiLike) else value
elif isinstance(idx, tuple) and Ellipsis in idx:
idx = tuple(slice(None) if i is Ellipsis else i for i in idx)
self.array[idx] = value.array if isinstance(value, CasadiLike) else value
else:
self.array[idx] = value.array if isinstance(value, CasadiLike) else value

def __getitem__(self, idx) -> "CasadiLike":
"""Overrides get item operator"""
return CasadiLike(self.array[idx])
if idx is Ellipsis:
# Handle the case where only Ellipsis is passed
return CasadiLike(self.array)
elif isinstance(idx, tuple) and Ellipsis in idx:
# Handle the case where Ellipsis is part of a tuple
idx = tuple(slice(None) if k is Ellipsis else k for k in idx)
return CasadiLike(self.array[idx])
else:
# For other cases, delegate to the CasADi object's __getitem__
return CasadiLike(self.array[idx])

@property
def T(self) -> "CasadiLike":
Expand Down
Loading