refactor `linalg`

docs/modules/linalg.md

@@ -5,6 +5,8 @@ label: linalg
 
 Provide `mat3x3`, `vec2`, `vec3` and `vec4` types.
 
+This classes adopt `torch`'s naming convention. Methods with `_` suffix will modify the instance itself.
+
 https://github.com/blueloveTH/pocketpy/blob/main/include/typings/linalg.pyi
 
 ```python
@@ -15,8 +17,6 @@ class vec2(_StructLike['vec2']):
     x: float
     y: float
 
-    def assign(self, other: vec2) -> None: ...
-
     def __init__(self, x: float, y: float) -> None: ...
     def __add__(self, other: vec2) -> vec2: ...
     def __sub__(self, other: vec2) -> vec2: ...
@@ -30,6 +30,10 @@ class vec2(_StructLike['vec2']):
     def normalize(self) -> vec2: ...
     def rotate(self, radians: float) -> vec2: ...
 
+    def copy_(self, other: vec2) -> None: ...
+    def normalize_(self) -> None: ...
+    def rotate_(self, radians: float) -> None: ...
+
     @staticmethod
     def angle(__from: vec2, __to: vec2) -> float:
         """Returns the angle in radians between vectors `from` and `to`.
@@ -49,8 +53,6 @@ class vec3(_StructLike['vec3']):
     y: float
     z: float
 
-    def assign(self, other: vec3) -> None: ...
-
     def __init__(self, x: float, y: float, z: float) -> None: ...
     def __add__(self, other: vec3) -> vec3: ...
     def __sub__(self, other: vec3) -> vec3: ...
@@ -63,14 +65,15 @@ class vec3(_StructLike['vec3']):
     def length_squared(self) -> float: ...
     def normalize(self) -> vec3: ...
 
+    def copy_(self, other: vec3) -> None: ...
+    def normalize_(self) -> None: ...
+
 class vec4(_StructLike['vec4']):
     x: float
     y: float
     z: float
     w: float
 
-    def assign(self, other: vec4) -> None: ...
-
     def __init__(self, x: float, y: float, z: float, w: float) -> None: ...
     def __add__(self, other: vec4) -> vec4: ...
     def __sub__(self, other: vec4) -> vec4: ...
@@ -82,6 +85,9 @@ class vec4(_StructLike['vec4']):
     def length_squared(self) -> float: ...
     def normalize(self) -> vec4: ...
 
+    def copy_(self, other: vec4) -> None: ...
+    def normalize_(self) -> None: ...
+
 class mat3x3(_StructLike['mat3x3']):
     _11: float
     _12: float
@@ -93,8 +99,6 @@ class mat3x3(_StructLike['mat3x3']):
     _32: float
     _33: float
 
-    def assign(self, other: mat3x3) -> None: ...
-
     @overload
     def __init__(self) -> None: ...
     @overload
@@ -102,11 +106,8 @@ class mat3x3(_StructLike['mat3x3']):
     @overload
     def __init__(self, a: list[float]): ...
 
-    def set_zeros(self) -> None: ...
-    def set_ones(self) -> None: ...
-    def set_identity(self) -> None: ...
-
     def determinant(self) -> float: ...
+    def invert(self) -> mat3x3: ...
     def transpose(self) -> mat3x3: ...
 
     def __getitem__(self, index: tuple[int, int]) -> float: ...
@@ -123,8 +124,11 @@ class mat3x3(_StructLike['mat3x3']):
     @overload
     def __matmul__(self, other: vec3) -> vec3: ...
 
-    def __imatmul__(self, other: mat3x3): ...
+    def matmul(self, other: mat3x3, out: mat3x3 = None) -> mat3x3 | None: ...
+
+    def copy_(self, other: mat3x3) -> None: ...
     def invert_(self) -> None: ...
+    def transpose_(self) -> None: ...
 
     @staticmethod
     def zeros() -> mat3x3: ...
@@ -136,15 +140,14 @@ class mat3x3(_StructLike['mat3x3']):
     # affine transformations
     @staticmethod
     def trs(t: vec2, r: float, s: vec2) -> mat3x3: ...
-    
-    def set_trs(self, t: vec2, r: float, s: vec2) -> None: ...
-
-    def is_affine(self) -> bool: ...
+    def copy_trs_(self, t: vec2, r: float, s: vec2) -> None: ...
 
     def _t(self) -> vec2: ...
     def _r(self) -> float: ...
     def _s(self) -> vec2: ...
 
+    def is_affine(self) -> bool: ...
+
     def transform_point(self, p: vec2) -> vec2: ...
     def transform_vector(self, v: vec2) -> vec2: ...
 

include/pocketpy/linalg.h

@@ -28,7 +28,9 @@ struct Vec2{
     float length() const { return sqrtf(x * x + y * y); }
     float length_squared() const { return x * x + y * y; }
     Vec2 normalize() const { float l = length(); return Vec2(x / l, y / l); }
-    NoReturn assign(const Vec2& v) { x = v.x; y = v.y; return {}; }
+    Vec2 rotate(float radian) const { float cr = cosf(radian), sr = sinf(radian); return Vec2(x * cr - y * sr, x * sr + y * cr); }
+    NoReturn normalize_() { float l = length(); x /= l; y /= l; return {}; }
+    NoReturn copy_(const Vec2& v) { x = v.x; y = v.y; return {}; }
 };
 
 struct Vec3{
@@ -53,7 +55,8 @@ struct Vec3{
     float length() const { return sqrtf(x * x + y * y + z * z); }
     float length_squared() const { return x * x + y * y + z * z; }
     Vec3 normalize() const { float l = length(); return Vec3(x / l, y / l, z / l); }
-    NoReturn assign(const Vec3& v) { x = v.x; y = v.y; z = v.z; return {}; }
+    NoReturn normalize_() { float l = length(); x /= l; y /= l; z /= l; return {}; }
+    NoReturn copy_(const Vec3& v) { x = v.x; y = v.y; z = v.z; return {}; }
 };
 
 struct Vec4{
@@ -77,7 +80,8 @@ struct Vec4{
     float length() const { return sqrtf(x * x + y * y + z * z + w * w); }
     float length_squared() const { return x * x + y * y + z * z + w * w; }
     Vec4 normalize() const { float l = length(); return Vec4(x / l, y / l, z / l, w / l); }
-    NoReturn assign(const Vec4& v) { x = v.x; y = v.y; z = v.z; w = v.w; return {}; }
+    NoReturn normalize_() { float l = length(); x /= l; y /= l; z /= l; w /= l; return {}; }
+    NoReturn copy_(const Vec4& v) { x = v.x; y = v.y; z = v.z; w = v.w; return {}; }
 };
 
 struct Mat3x3{    
@@ -95,10 +99,6 @@ struct Mat3x3{
     Mat3x3(float, float, float, float, float, float, float, float, float);
     Mat3x3(const Mat3x3& other) = default;
 
-    void set_zeros();
-    void set_ones();
-    void set_identity();
-
     static Mat3x3 zeros();
     static Mat3x3 ones();
     static Mat3x3 identity();

include/typings/linalg.pyi

@@ -5,8 +5,6 @@ class vec2(_StructLike['vec2']):
     x: float
     y: float
 
-    def assign(self, other: vec2) -> None: ...
-
     def __init__(self, x: float, y: float) -> None: ...
     def __add__(self, other: vec2) -> vec2: ...
     def __sub__(self, other: vec2) -> vec2: ...
@@ -20,6 +18,10 @@ class vec2(_StructLike['vec2']):
     def normalize(self) -> vec2: ...
     def rotate(self, radians: float) -> vec2: ...
 
+    def copy_(self, other: vec2) -> None: ...
+    def normalize_(self) -> None: ...
+    def rotate_(self, radians: float) -> None: ...
+
     @staticmethod
     def angle(__from: vec2, __to: vec2) -> float:
         """Returns the angle in radians between vectors `from` and `to`.
@@ -39,8 +41,6 @@ class vec3(_StructLike['vec3']):
     y: float
     z: float
 
-    def assign(self, other: vec3) -> None: ...
-
     def __init__(self, x: float, y: float, z: float) -> None: ...
     def __add__(self, other: vec3) -> vec3: ...
     def __sub__(self, other: vec3) -> vec3: ...
@@ -53,14 +53,15 @@ class vec3(_StructLike['vec3']):
     def length_squared(self) -> float: ...
     def normalize(self) -> vec3: ...
 
+    def copy_(self, other: vec3) -> None: ...
+    def normalize_(self) -> None: ...
+
 class vec4(_StructLike['vec4']):
     x: float
     y: float
     z: float
     w: float
 
-    def assign(self, other: vec4) -> None: ...
-
     def __init__(self, x: float, y: float, z: float, w: float) -> None: ...
     def __add__(self, other: vec4) -> vec4: ...
     def __sub__(self, other: vec4) -> vec4: ...
@@ -72,6 +73,9 @@ class vec4(_StructLike['vec4']):
     def length_squared(self) -> float: ...
     def normalize(self) -> vec4: ...
 
+    def copy_(self, other: vec4) -> None: ...
+    def normalize_(self) -> None: ...
+
 class mat3x3(_StructLike['mat3x3']):
     _11: float
     _12: float
@@ -83,8 +87,6 @@ class mat3x3(_StructLike['mat3x3']):
     _32: float
     _33: float
 
-    def assign(self, other: mat3x3) -> None: ...
-
     @overload
     def __init__(self) -> None: ...
     @overload
@@ -92,11 +94,8 @@ class mat3x3(_StructLike['mat3x3']):
     @overload
     def __init__(self, a: list[float]): ...
 
-    def set_zeros(self) -> None: ...
-    def set_ones(self) -> None: ...
-    def set_identity(self) -> None: ...
-
     def determinant(self) -> float: ...
+    def invert(self) -> mat3x3: ...
     def transpose(self) -> mat3x3: ...
 
     def __getitem__(self, index: tuple[int, int]) -> float: ...
@@ -113,8 +112,11 @@ class mat3x3(_StructLike['mat3x3']):
     @overload
     def __matmul__(self, other: vec3) -> vec3: ...
 
-    def __imatmul__(self, other: mat3x3): ...
+    def matmul(self, other: mat3x3, out: mat3x3 = None) -> mat3x3 | None: ...
+
+    def copy_(self, other: mat3x3) -> None: ...
     def invert_(self) -> None: ...
+    def transpose_(self) -> None: ...
 
     @staticmethod
     def zeros() -> mat3x3: ...
@@ -126,15 +128,14 @@ class mat3x3(_StructLike['mat3x3']):
     # affine transformations
     @staticmethod
     def trs(t: vec2, r: float, s: vec2) -> mat3x3: ...
-    
-    def set_trs(self, t: vec2, r: float, s: vec2) -> None: ...
-
-    def is_affine(self) -> bool: ...
+    def copy_trs_(self, t: vec2, r: float, s: vec2) -> None: ...
 
     def _t(self) -> vec2: ...
     def _r(self) -> float: ...
     def _s(self) -> vec2: ...
 
+    def is_affine(self) -> bool: ...
+
     def transform_point(self, p: vec2) -> vec2: ...
     def transform_vector(self, v: vec2) -> vec2: ...
 

src/linalg.cpp

@@ -127,10 +127,14 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
         vm->bind_method<1>(type, "rotate", [](VM* vm, ArgsView args){
             Vec2 self = _CAST(PyVec2&, args[0]);
             float radian = CAST(f64, args[1]);
-            float cr = cosf(radian);
-            float sr = sinf(radian);
-            Vec2 res(self.x * cr - self.y * sr, self.x * sr + self.y * cr);
-            return VAR_T(PyVec2, res);
+            return VAR_T(PyVec2, self.rotate(radian));
+        });
+
+        vm->bind_method<0>(type, "rotate_", [](VM* vm, ArgsView args){
+            PyVec2& self = _CAST(PyVec2&, args[0]);
+            float radian = CAST(f64, args[1]);
+            self = self.rotate(radian);
+            return vm->None;
         });
 
         PY_FIELD(PyVec2, "x", _, x)
@@ -143,10 +147,11 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
         BIND_VEC_FLOAT_OP(2, __truediv__, /)
         BIND_VEC_FUNCTION_1(2, dot)
         BIND_VEC_FUNCTION_1(2, cross)
-        BIND_VEC_FUNCTION_1(2, assign)
+        BIND_VEC_FUNCTION_1(2, copy_)
         BIND_VEC_FUNCTION_0(2, length)
         BIND_VEC_FUNCTION_0(2, length_squared)
         BIND_VEC_FUNCTION_0(2, normalize)
+        BIND_VEC_FUNCTION_0(2, normalize_)
     }
 
     void PyVec3::_register(VM* vm, PyObject* mod, PyObject* type){
@@ -178,10 +183,11 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
         BIND_VEC_FLOAT_OP(3, __truediv__, /)
         BIND_VEC_FUNCTION_1(3, dot)
         BIND_VEC_FUNCTION_1(3, cross)
-        BIND_VEC_FUNCTION_1(3, assign)
+        BIND_VEC_FUNCTION_1(3, copy_)
         BIND_VEC_FUNCTION_0(3, length)
         BIND_VEC_FUNCTION_0(3, length_squared)
         BIND_VEC_FUNCTION_0(3, normalize)
+        BIND_VEC_FUNCTION_0(3, normalize_)
     }
 
     void PyVec4::_register(VM* vm, PyObject* mod, PyObject* type){
@@ -214,10 +220,11 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
         BIND_VEC_FLOAT_OP(4, __rmul__, *)
         BIND_VEC_FLOAT_OP(4, __truediv__, /)
         BIND_VEC_FUNCTION_1(4, dot)
-        BIND_VEC_FUNCTION_1(4, assign)
+        BIND_VEC_FUNCTION_1(4, copy_)
         BIND_VEC_FUNCTION_0(4, length)
         BIND_VEC_FUNCTION_0(4, length_squared)
         BIND_VEC_FUNCTION_0(4, normalize)
+        BIND_VEC_FUNCTION_0(4, normalize_)
     }
 
 #undef BIND_VEC_VEC_OP
@@ -246,17 +253,13 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
             return vm->None;
         });
 
-        vm->bind_method<1>(type, "assign", [](VM* vm, ArgsView args){
+        vm->bind_method<1>(type, "copy_", [](VM* vm, ArgsView args){
             PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
             const PyMat3x3& other = CAST(PyMat3x3&, args[1]);
             self = other;
             return vm->None;
         });
 
-        vm->bind_method<0>(type, "set_zeros", PK_ACTION(PK_OBJ_GET(PyMat3x3, args[0]).set_zeros()));
-        vm->bind_method<0>(type, "set_ones", PK_ACTION(PK_OBJ_GET(PyMat3x3, args[0]).set_ones()));
-        vm->bind_method<0>(type, "set_identity", PK_ACTION(PK_OBJ_GET(PyMat3x3, args[0]).set_identity()));
-
         vm->bind__repr__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* obj){
             PyMat3x3& self = _CAST(PyMat3x3&, obj);
             std::stringstream ss;
@@ -352,12 +355,16 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
             return vm->NotImplemented;
         });
 
-        vm->bind_method<1>(type, "__imatmul__", [](VM* vm, ArgsView args){
-            PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
-            vm->check_non_tagged_type(args[1], PyMat3x3::_type(vm));
-            const PyMat3x3& other = _CAST(PyMat3x3&, args[1]);
-            self = self.matmul(other);
-            return args[0];
+        vm->bind(type, "matmul(self, other: mat3x3, out: mat3x3 = None)", [](VM* vm, ArgsView args){
+            const PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
+            const PyMat3x3& other = CAST(PyMat3x3&, args[1]);
+            if(args[2] == vm->None){
+                return VAR_T(PyMat3x3, self.matmul(other));
+            }else{
+                PyMat3x3& out = _CAST(PyMat3x3&, args[2]);
+                out = self.matmul(other);
+                return vm->None;
+            }
         });
 
         vm->bind_method<0>(type, "determinant", [](VM* vm, ArgsView args){
@@ -378,6 +385,14 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
             return VAR_T(PyMat3x3, ret);
         });
 
+        vm->bind_method<0>(type, "invert", [](VM* vm, ArgsView args){
+            PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
+            Mat3x3 ret;
+            bool ok = self.inverse(ret);
+            if(!ok) vm->ValueError("matrix is not invertible");
+            return VAR_T(PyMat3x3, ret);
+        });
+
         vm->bind_method<0>(type, "invert_", [](VM* vm, ArgsView args){
             PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
             Mat3x3 ret;
@@ -387,6 +402,12 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
             return vm->None;
         });
 
+        vm->bind_method<0>(type, "transpose_", [](VM* vm, ArgsView args){
+            PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
+            self = self.transpose();
+            return vm->None;
+        });
+
         // @staticmethod
         vm->bind(type, "zeros()", [](VM* vm, ArgsView args){
             PK_UNUSED(args);
@@ -414,7 +435,7 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, PyVec2& currentVelocity, float
             return VAR_T(PyMat3x3, Mat3x3::trs(t, r, s));
         }, {}, BindType::STATICMETHOD);
 
-        vm->bind(type, "set_trs(self, t: vec2, r: float, s: vec2)", [](VM* vm, ArgsView args){
+        vm->bind(type, "copy_trs_(self, t: vec2, r: float, s: vec2)", [](VM* vm, ArgsView args){
             PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
             Vec2 t = CAST(Vec2, args[1]);
             f64 r = CAST_F(args[2]);
@@ -483,10 +504,6 @@ void add_module_linalg(VM* vm){
         , _21(_21), _22(_22), _23(_23)
         , _31(_31), _32(_32), _33(_33) {}
 
-    void Mat3x3::set_zeros(){ for (int i=0; i<9; ++i) v[i] = 0.0f; }
-    void Mat3x3::set_ones(){ for (int i=0; i<9; ++i) v[i] = 1.0f; }
-    void Mat3x3::set_identity(){ set_zeros(); _11 = _22 = _33 = 1.0f; }
-
     Mat3x3 Mat3x3::zeros(){
         return Mat3x3(0, 0, 0, 0, 0, 0, 0, 0, 0);
     }

tests/80_linalg.py

@@ -221,23 +221,6 @@ test_mat_copy = test_mat.copy()
 assert test_mat is not test_mat_copy
 assert test_mat == test_mat_copy
 
-# test setzeros
-test_mat_copy = test_mat.copy()
-test_mat_copy.set_zeros()
-assert test_mat_copy == mat3x3.zeros()
-
-# test set_ones
-test_mat_copy = test_mat.copy()
-test_mat_copy.set_ones()
-assert test_mat_copy == mat3x3.ones()
-
-# test set_identity
-test_mat_copy = test_mat.copy()
-test_mat_copy.set_identity()
-assert test_mat_copy == mat3x3([1, 0, 0,
-                                0, 1, 0,
-                                0, 0, 1])
-
 # test __getitem__
 for i, element in enumerate([getattr(test_mat, e) for e in element_name_list]):
     assert test_mat[int(i/3), i%3] == element
@@ -276,8 +259,7 @@ except:
 
 # test __add__
 test_mat_copy = test_mat.copy()
-ones = mat3x3()
-ones.set_ones()
+ones = mat3x3.ones()
 result_mat = test_mat_copy.__add__(ones)
 correct_result_mat = test_mat_copy.copy()
 for i in range(3):
@@ -287,8 +269,7 @@ assert result_mat == correct_result_mat
 
 # test __sub__
 test_mat_copy = test_mat.copy()
-ones = mat3x3()
-ones.set_ones()
+ones = mat3x3.ones()
 result_mat = test_mat_copy.__sub__(ones)
 correct_result_mat = test_mat_copy.copy()
 for i in range(3):
@@ -318,9 +299,6 @@ for i in range(3):
         correct_result_mat[i, j] = sum([e1*e2 for e1, e2 in zip(get_row(test_mat_copy, i), get_col(test_mat_copy_2, j))])
 assert result_mat == correct_result_mat
 
-test_mat_copy.__imatmul__(test_mat_copy_2)
-assert test_mat_copy == correct_result_mat
-
 # test determinant
 test_mat_copy = test_mat.copy()
 test_mat_copy.determinant()
@@ -408,7 +386,7 @@ radian = random.uniform(-10*math.pi, 10*math.pi)
 assert mat_to_str_list(mat3x3.trs(test_vec2_copy, radian, test_vec2_2_copy)) == mat_list_to_str_list(trs(test_vec2_list, radian, test_vec2_2_list))
 
 a = mat3x3.zeros()
-a.set_trs(test_vec2_copy, radian, test_vec2_2_copy)
+a.copy_trs_(test_vec2_copy, radian, test_vec2_2_copy)
 assert a == mat3x3.trs(test_vec2_copy, radian, test_vec2_2_copy)
 
 # test is_affine
@@ -474,18 +452,18 @@ assert mymat3x3().f()
 
 # test assign
 a = vec2(1, 2)
-assert a.assign(vec2(3, 4)) is None
+assert a.copy_(vec2(3, 4)) is None
 assert a == vec2(3, 4)
 
 b = vec3(1, 2, 3)
-assert b.assign(vec3(4, 5, 6)) is None
+assert b.copy_(vec3(4, 5, 6)) is None
 assert b == vec3(4, 5, 6)
 
 c = vec4(1, 2, 3, 4)
-assert c.assign(vec4(5, 6, 7, 8)) is None
+assert c.copy_(vec4(5, 6, 7, 8)) is None
 assert c == vec4(5, 6, 7, 8)
 
 d = mat3x3.identity()
-assert d.assign(mat3x3.zeros()) is None
+assert d.copy_(mat3x3.zeros()) is None
 assert d == mat3x3.zeros()