linalg.h 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676
  1. #pragma once
  2. #include <cmath>
  3. #include "cffi.h"
  4. namespace pkpy{
  5. static constexpr float kEpsilon = 1e-4f;
  6. inline static bool isclose(float a, float b){ return fabsf(a - b) < kEpsilon; }
  7. struct Vec2{
  8. float x, y;
  9. Vec2() : x(0.0f), y(0.0f) {}
  10. Vec2(float x, float y) : x(x), y(y) {}
  11. Vec2(const Vec2& v) : x(v.x), y(v.y) {}
  12. Vec2 operator+(const Vec2& v) const { return Vec2(x + v.x, y + v.y); }
  13. Vec2& operator+=(const Vec2& v) { x += v.x; y += v.y; return *this; }
  14. Vec2 operator-(const Vec2& v) const { return Vec2(x - v.x, y - v.y); }
  15. Vec2& operator-=(const Vec2& v) { x -= v.x; y -= v.y; return *this; }
  16. Vec2 operator*(float s) const { return Vec2(x * s, y * s); }
  17. Vec2& operator*=(float s) { x *= s; y *= s; return *this; }
  18. Vec2 operator/(float s) const { return Vec2(x / s, y / s); }
  19. Vec2& operator/=(float s) { x /= s; y /= s; return *this; }
  20. Vec2 operator-() const { return Vec2(-x, -y); }
  21. bool operator==(const Vec2& v) const { return isclose(x, v.x) && isclose(y, v.y); }
  22. bool operator!=(const Vec2& v) const { return !isclose(x, v.x) || !isclose(y, v.y); }
  23. float dot(const Vec2& v) const { return x * v.x + y * v.y; }
  24. float cross(const Vec2& v) const { return x * v.y - y * v.x; }
  25. float length() const { return sqrtf(x * x + y * y); }
  26. float length_squared() const { return x * x + y * y; }
  27. Vec2 normalize() const { float l = length(); return Vec2(x / l, y / l); }
  28. };
  29. struct Vec3{
  30. float x, y, z;
  31. Vec3() : x(0.0f), y(0.0f), z(0.0f) {}
  32. Vec3(float x, float y, float z) : x(x), y(y), z(z) {}
  33. Vec3(const Vec3& v) : x(v.x), y(v.y), z(v.z) {}
  34. Vec3 operator+(const Vec3& v) const { return Vec3(x + v.x, y + v.y, z + v.z); }
  35. Vec3& operator+=(const Vec3& v) { x += v.x; y += v.y; z += v.z; return *this; }
  36. Vec3 operator-(const Vec3& v) const { return Vec3(x - v.x, y - v.y, z - v.z); }
  37. Vec3& operator-=(const Vec3& v) { x -= v.x; y -= v.y; z -= v.z; return *this; }
  38. Vec3 operator*(float s) const { return Vec3(x * s, y * s, z * s); }
  39. Vec3& operator*=(float s) { x *= s; y *= s; z *= s; return *this; }
  40. Vec3 operator/(float s) const { return Vec3(x / s, y / s, z / s); }
  41. Vec3& operator/=(float s) { x /= s; y /= s; z /= s; return *this; }
  42. Vec3 operator-() const { return Vec3(-x, -y, -z); }
  43. bool operator==(const Vec3& v) const { return isclose(x, v.x) && isclose(y, v.y) && isclose(z, v.z); }
  44. bool operator!=(const Vec3& v) const { return !isclose(x, v.x) || !isclose(y, v.y) || !isclose(z, v.z); }
  45. float dot(const Vec3& v) const { return x * v.x + y * v.y + z * v.z; }
  46. Vec3 cross(const Vec3& v) const { return Vec3(y * v.z - z * v.y, z * v.x - x * v.z, x * v.y - y * v.x); }
  47. float length() const { return sqrtf(x * x + y * y + z * z); }
  48. float length_squared() const { return x * x + y * y + z * z; }
  49. Vec3 normalize() const { float l = length(); return Vec3(x / l, y / l, z / l); }
  50. };
  51. struct Mat3x3{
  52. union {
  53. struct {
  54. float _11, _12, _13;
  55. float _21, _22, _23;
  56. float _31, _32, _33;
  57. };
  58. float m[3][3];
  59. float v[9];
  60. };
  61. Mat3x3() {}
  62. Mat3x3(float _11, float _12, float _13,
  63. float _21, float _22, float _23,
  64. float _31, float _32, float _33)
  65. : _11(_11), _12(_12), _13(_13)
  66. , _21(_21), _22(_22), _23(_23)
  67. , _31(_31), _32(_32), _33(_33) {}
  68. void set_zeros(){ for (int i=0; i<9; ++i) v[i] = 0.0f; }
  69. void set_ones(){ for (int i=0; i<9; ++i) v[i] = 1.0f; }
  70. void set_identity(){ set_zeros(); _11 = _22 = _33 = 1.0f; }
  71. static Mat3x3 zeros(){
  72. static Mat3x3 ret(0, 0, 0, 0, 0, 0, 0, 0, 0);
  73. return ret;
  74. }
  75. static Mat3x3 ones(){
  76. static Mat3x3 ret(1, 1, 1, 1, 1, 1, 1, 1, 1);
  77. return ret;
  78. }
  79. static Mat3x3 identity(){
  80. static Mat3x3 ret(1, 0, 0, 0, 1, 0, 0, 0, 1);
  81. return ret;
  82. }
  83. Mat3x3 operator+(const Mat3x3& other) const{
  84. Mat3x3 ret;
  85. for (int i=0; i<9; ++i) ret.v[i] = v[i] + other.v[i];
  86. return ret;
  87. }
  88. Mat3x3 operator-(const Mat3x3& other) const{
  89. Mat3x3 ret;
  90. for (int i=0; i<9; ++i) ret.v[i] = v[i] - other.v[i];
  91. return ret;
  92. }
  93. Mat3x3 operator*(float scalar) const{
  94. Mat3x3 ret;
  95. for (int i=0; i<9; ++i) ret.v[i] = v[i] * scalar;
  96. return ret;
  97. }
  98. Mat3x3 operator/(float scalar) const{
  99. Mat3x3 ret;
  100. for (int i=0; i<9; ++i) ret.v[i] = v[i] / scalar;
  101. return ret;
  102. }
  103. Mat3x3& operator+=(const Mat3x3& other){
  104. for (int i=0; i<9; ++i) v[i] += other.v[i];
  105. return *this;
  106. }
  107. Mat3x3& operator-=(const Mat3x3& other){
  108. for (int i=0; i<9; ++i) v[i] -= other.v[i];
  109. return *this;
  110. }
  111. Mat3x3& operator*=(float scalar){
  112. for (int i=0; i<9; ++i) v[i] *= scalar;
  113. return *this;
  114. }
  115. Mat3x3& operator/=(float scalar){
  116. for (int i=0; i<9; ++i) v[i] /= scalar;
  117. return *this;
  118. }
  119. Mat3x3 matmul(const Mat3x3& other) const{
  120. Mat3x3 ret;
  121. ret._11 = _11 * other._11 + _12 * other._21 + _13 * other._31;
  122. ret._12 = _11 * other._12 + _12 * other._22 + _13 * other._32;
  123. ret._13 = _11 * other._13 + _12 * other._23 + _13 * other._33;
  124. ret._21 = _21 * other._11 + _22 * other._21 + _23 * other._31;
  125. ret._22 = _21 * other._12 + _22 * other._22 + _23 * other._32;
  126. ret._23 = _21 * other._13 + _22 * other._23 + _23 * other._33;
  127. ret._31 = _31 * other._11 + _32 * other._21 + _33 * other._31;
  128. ret._32 = _31 * other._12 + _32 * other._22 + _33 * other._32;
  129. ret._33 = _31 * other._13 + _32 * other._23 + _33 * other._33;
  130. return ret;
  131. }
  132. Vec3 matmul(const Vec3& other) const{
  133. Vec3 ret;
  134. ret.x = _11 * other.x + _12 * other.y + _13 * other.z;
  135. ret.y = _21 * other.x + _22 * other.y + _23 * other.z;
  136. ret.z = _31 * other.x + _32 * other.y + _33 * other.z;
  137. return ret;
  138. }
  139. bool operator==(const Mat3x3& other) const{
  140. for (int i=0; i<9; ++i){
  141. if (!isclose(v[i], other.v[i])) return false;
  142. }
  143. return true;
  144. }
  145. bool operator!=(const Mat3x3& other) const{
  146. for (int i=0; i<9; ++i){
  147. if (!isclose(v[i], other.v[i])) return true;
  148. }
  149. return false;
  150. }
  151. float determinant() const{
  152. return _11 * _22 * _33 + _12 * _23 * _31 + _13 * _21 * _32
  153. - _11 * _23 * _32 - _12 * _21 * _33 - _13 * _22 * _31;
  154. }
  155. Mat3x3 transpose() const{
  156. Mat3x3 ret;
  157. ret._11 = _11; ret._12 = _21; ret._13 = _31;
  158. ret._21 = _12; ret._22 = _22; ret._23 = _32;
  159. ret._31 = _13; ret._32 = _23; ret._33 = _33;
  160. return ret;
  161. }
  162. bool inverse(Mat3x3& ret) const{
  163. float det = determinant();
  164. if (fabsf(det) < kEpsilon) return false;
  165. float inv_det = 1.0f / det;
  166. ret._11 = (_22 * _33 - _23 * _32) * inv_det;
  167. ret._12 = (_13 * _32 - _12 * _33) * inv_det;
  168. ret._13 = (_12 * _23 - _13 * _22) * inv_det;
  169. ret._21 = (_23 * _31 - _21 * _33) * inv_det;
  170. ret._22 = (_11 * _33 - _13 * _31) * inv_det;
  171. ret._23 = (_13 * _21 - _11 * _23) * inv_det;
  172. ret._31 = (_21 * _32 - _22 * _31) * inv_det;
  173. ret._32 = (_12 * _31 - _11 * _32) * inv_det;
  174. ret._33 = (_11 * _22 - _12 * _21) * inv_det;
  175. return true;
  176. }
  177. /*************** affine transformations ***************/
  178. static Mat3x3 trs(Vec2 t, float radian, Vec2 s){
  179. float cr = cosf(radian);
  180. float sr = sinf(radian);
  181. return Mat3x3(s.x * cr, -s.y * sr, t.x,
  182. s.x * sr, s.y * cr, t.y,
  183. 0.0f, 0.0f, 1.0f);
  184. }
  185. bool is_affine() const{
  186. float det = _11 * _22 - _12 * _21;
  187. if(fabsf(det) < kEpsilon) return false;
  188. return _31 == 0.0f && _32 == 0.0f && _33 == 1.0f;
  189. }
  190. Mat3x3 inverse_affine() const{
  191. Mat3x3 ret;
  192. float det = _11 * _22 - _12 * _21;
  193. float inv_det = 1.0f / det;
  194. ret._11 = _22 * inv_det;
  195. ret._12 = -_12 * inv_det;
  196. ret._13 = (_12 * _23 - _13 * _22) * inv_det;
  197. ret._21 = -_21 * inv_det;
  198. ret._22 = _11 * inv_det;
  199. ret._23 = (_13 * _21 - _11 * _23) * inv_det;
  200. ret._31 = 0.0f;
  201. ret._32 = 0.0f;
  202. ret._33 = 1.0f;
  203. return ret;
  204. }
  205. Mat3x3 matmul_affine(const Mat3x3& other) const{
  206. Mat3x3 ret;
  207. ret._11 = _11 * other._11 + _12 * other._21;
  208. ret._12 = _11 * other._12 + _12 * other._22;
  209. ret._13 = _11 * other._13 + _12 * other._23 + _13;
  210. ret._21 = _21 * other._11 + _22 * other._21;
  211. ret._22 = _21 * other._12 + _22 * other._22;
  212. ret._23 = _21 * other._13 + _22 * other._23 + _23;
  213. ret._31 = 0.0f;
  214. ret._32 = 0.0f;
  215. ret._33 = 1.0f;
  216. return ret;
  217. }
  218. Vec2 translation() const { return Vec2(_13, _23); }
  219. float rotation() const { return atan2f(_21, _11); }
  220. Vec2 scale() const {
  221. return Vec2(
  222. sqrtf(_11 * _11 + _21 * _21),
  223. sqrtf(_12 * _12 + _22 * _22)
  224. );
  225. }
  226. Vec2 transform_point(Vec2 v) const {
  227. return Vec2(_11 * v.x + _12 * v.y + _13, _21 * v.x + _22 * v.y + _23);
  228. }
  229. Vec2 transform_vector(Vec2 v) const {
  230. return Vec2(_11 * v.x + _12 * v.y, _21 * v.x + _22 * v.y);
  231. }
  232. };
  233. struct PyVec2;
  234. struct PyVec3;
  235. struct PyMat3x3;
  236. PyObject* py_var(VM*, Vec2);
  237. PyObject* py_var(VM*, const PyVec2&);
  238. PyObject* py_var(VM*, Vec3);
  239. PyObject* py_var(VM*, const PyVec3&);
  240. PyObject* py_var(VM*, const Mat3x3&);
  241. PyObject* py_var(VM*, const PyMat3x3&);
  242. #define BIND_VEC_VEC_OP(D, name, op) \
  243. vm->bind_method<1>(type, #name, [](VM* vm, ArgsView args){ \
  244. PyVec##D& self = _CAST(PyVec##D&, args[0]); \
  245. PyVec##D& other = CAST(PyVec##D&, args[1]); \
  246. return VAR(self op other); \
  247. });
  248. #define BIND_VEC_FLOAT_OP(D, name, op) \
  249. vm->bind_method<1>(type, #name, [](VM* vm, ArgsView args){ \
  250. PyVec##D& self = _CAST(PyVec##D&, args[0]); \
  251. f64 other = vm->num_to_float(args[1]); \
  252. return VAR(self op other); \
  253. });
  254. #define BIND_VEC_FUNCTION_0(D, name) \
  255. vm->bind_method<0>(type, #name, [](VM* vm, ArgsView args){ \
  256. PyVec##D& self = _CAST(PyVec##D&, args[0]); \
  257. return VAR(self.name()); \
  258. });
  259. #define BIND_VEC_FUNCTION_1(D, name) \
  260. vm->bind_method<0>(type, #name, [](VM* vm, ArgsView args){ \
  261. PyVec##D& self = _CAST(PyVec##D&, args[0]); \
  262. PyVec##D& other = CAST(PyVec##D&, args[1]); \
  263. return VAR(self.name(other)); \
  264. });
  265. #define BIND_VEC_FIELD(D, name) \
  266. type->attr().set(#name, vm->property([](VM* vm, ArgsView args){ \
  267. PyVec##D& self = _CAST(PyVec##D&, args[0]); \
  268. return VAR(self.name); \
  269. }, [](VM* vm, ArgsView args){ \
  270. PyVec##D& self = _CAST(PyVec##D&, args[0]); \
  271. self.name = vm->num_to_float(args[1]); \
  272. return vm->None; \
  273. }));
  274. struct PyVec2: Vec2 {
  275. PY_CLASS(PyVec2, linalg, vec2)
  276. PyVec2() : Vec2() {}
  277. PyVec2(const Vec2& v) : Vec2(v) {}
  278. PyVec2(const PyVec2& v) : Vec2(v) {}
  279. static void _register(VM* vm, PyObject* mod, PyObject* type){
  280. vm->bind_constructor<3>(type, [](VM* vm, ArgsView args){
  281. float x = CAST_F(args[1]);
  282. float y = CAST_F(args[2]);
  283. return VAR(Vec2(x, y));
  284. });
  285. vm->bind__repr__(OBJ_GET(Type, type), [](VM* vm, PyObject* obj){
  286. PyVec2& self = _CAST(PyVec2&, obj);
  287. std::stringstream ss;
  288. ss << "vec2(" << self.x << ", " << self.y << ")";
  289. return VAR(ss.str());
  290. });
  291. vm->bind_method<0>(type, "copy", [](VM* vm, ArgsView args){
  292. PyVec2& self = _CAST(PyVec2&, args[0]);
  293. return VAR_T(PyVec2, self);
  294. });
  295. vm->bind_method<1>(type, "rotate", [](VM* vm, ArgsView args){
  296. Vec2 self = _CAST(PyVec2&, args[0]);
  297. float radian = vm->num_to_float(args[1]);
  298. float cr = cosf(radian);
  299. float sr = sinf(radian);
  300. Mat3x3 rotate(cr, -sr, 0.0f,
  301. sr, cr, 0.0f,
  302. 0.0f, 0.0f, 1.0f);
  303. self = rotate.transform_vector(self);
  304. return VAR(self);
  305. });
  306. BIND_VEC_VEC_OP(2, __add__, +)
  307. BIND_VEC_VEC_OP(2, __sub__, -)
  308. BIND_VEC_FLOAT_OP(2, __mul__, *)
  309. BIND_VEC_FLOAT_OP(2, __truediv__, /)
  310. BIND_VEC_VEC_OP(2, __eq__, ==)
  311. BIND_VEC_FIELD(2, x)
  312. BIND_VEC_FIELD(2, y)
  313. BIND_VEC_FUNCTION_1(2, dot)
  314. BIND_VEC_FUNCTION_1(2, cross)
  315. BIND_VEC_FUNCTION_0(2, length)
  316. BIND_VEC_FUNCTION_0(2, length_squared)
  317. BIND_VEC_FUNCTION_0(2, normalize)
  318. }
  319. };
  320. struct PyVec3: Vec3 {
  321. PY_CLASS(PyVec3, linalg, vec3)
  322. PyVec3() : Vec3() {}
  323. PyVec3(const Vec3& v) : Vec3(v) {}
  324. PyVec3(const PyVec3& v) : Vec3(v) {}
  325. static void _register(VM* vm, PyObject* mod, PyObject* type){
  326. vm->bind_constructor<4>(type, [](VM* vm, ArgsView args){
  327. float x = CAST_F(args[1]);
  328. float y = CAST_F(args[2]);
  329. float z = CAST_F(args[3]);
  330. return VAR(Vec3(x, y, z));
  331. });
  332. vm->bind__repr__(OBJ_GET(Type, type), [](VM* vm, PyObject* obj){
  333. PyVec3& self = _CAST(PyVec3&, obj);
  334. std::stringstream ss;
  335. ss << "vec3(" << self.x << ", " << self.y << ", " << self.z << ")";
  336. return VAR(ss.str());
  337. });
  338. vm->bind_method<0>(type, "copy", [](VM* vm, ArgsView args){
  339. PyVec3& self = _CAST(PyVec3&, args[0]);
  340. return VAR_T(PyVec3, self);
  341. });
  342. BIND_VEC_VEC_OP(3, __add__, +)
  343. BIND_VEC_VEC_OP(3, __sub__, -)
  344. BIND_VEC_FLOAT_OP(3, __mul__, *)
  345. BIND_VEC_FLOAT_OP(3, __truediv__, /)
  346. BIND_VEC_VEC_OP(3, __eq__, ==)
  347. BIND_VEC_FIELD(3, x)
  348. BIND_VEC_FIELD(3, y)
  349. BIND_VEC_FIELD(3, z)
  350. BIND_VEC_FUNCTION_1(3, dot)
  351. BIND_VEC_FUNCTION_1(3, cross)
  352. BIND_VEC_FUNCTION_0(3, length)
  353. BIND_VEC_FUNCTION_0(3, length_squared)
  354. BIND_VEC_FUNCTION_0(3, normalize)
  355. }
  356. };
  357. struct PyMat3x3: Mat3x3{
  358. PY_CLASS(PyMat3x3, linalg, mat3x3)
  359. PyMat3x3(): Mat3x3(){}
  360. PyMat3x3(const Mat3x3& other): Mat3x3(other){}
  361. PyMat3x3(const PyMat3x3& other): Mat3x3(other){}
  362. static void _register(VM* vm, PyObject* mod, PyObject* type){
  363. vm->bind_constructor<-1>(type, [](VM* vm, ArgsView args){
  364. if(args.size() == 1+0) return VAR_T(PyMat3x3, Mat3x3::zeros());
  365. if(args.size() == 1+9){
  366. Mat3x3 mat;
  367. for(int i=0; i<9; i++) mat.v[i] = CAST_F(args[1+i]);
  368. return VAR_T(PyMat3x3, mat);
  369. }
  370. if(args.size() == 1+1){
  371. List& a = CAST(List&, args[1]);
  372. if(a.size() != 3) vm->ValueError("Mat3x3.__new__ takes 3x3 list");
  373. Mat3x3 mat;
  374. for(int i=0; i<3; i++){
  375. List& b = CAST(List&, a[i]);
  376. if(b.size() != 3) vm->ValueError("Mat3x3.__new__ takes 3x3 list");
  377. for(int j=0; j<3; j++){
  378. mat.m[i][j] = CAST_F(b[j]);
  379. }
  380. }
  381. return VAR_T(PyMat3x3, mat);
  382. }
  383. vm->TypeError("Mat3x3.__new__ takes 0 or 1 arguments");
  384. return vm->None;
  385. });
  386. #define METHOD_PROXY_NONE(name) \
  387. vm->bind_method<0>(type, #name, [](VM* vm, ArgsView args){ \
  388. PyMat3x3& self = _CAST(PyMat3x3&, args[0]); \
  389. self.name(); \
  390. return vm->None; \
  391. });
  392. METHOD_PROXY_NONE(set_zeros)
  393. METHOD_PROXY_NONE(set_ones)
  394. METHOD_PROXY_NONE(set_identity)
  395. #undef METHOD_PROXY_NONE
  396. vm->bind__repr__(OBJ_GET(Type, type), [](VM* vm, PyObject* obj){
  397. PyMat3x3& self = _CAST(PyMat3x3&, obj);
  398. std::stringstream ss;
  399. ss << std::fixed << std::setprecision(4);
  400. ss << "mat3x3([[" << self._11 << ", " << self._12 << ", " << self._13 << "],\n";
  401. ss << " [" << self._21 << ", " << self._22 << ", " << self._23 << "],\n";
  402. ss << " [" << self._31 << ", " << self._32 << ", " << self._33 << "]])";
  403. return VAR(ss.str());
  404. });
  405. vm->bind_method<0>(type, "copy", [](VM* vm, ArgsView args){
  406. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  407. return VAR_T(PyMat3x3, self);
  408. });
  409. vm->bind__getitem__(OBJ_GET(Type, type), [](VM* vm, PyObject* obj, PyObject* index){
  410. PyMat3x3& self = _CAST(PyMat3x3&, obj);
  411. Tuple& t = CAST(Tuple&, index);
  412. if(t.size() != 2){
  413. vm->TypeError("Mat3x3.__getitem__ takes a tuple of 2 integers");
  414. return vm->None;
  415. }
  416. i64 i = CAST(i64, t[0]);
  417. i64 j = CAST(i64, t[1]);
  418. if(i < 0 || i >= 3 || j < 0 || j >= 3){
  419. vm->IndexError("index out of range");
  420. return vm->None;
  421. }
  422. return VAR(self.m[i][j]);
  423. });
  424. vm->bind_method<2>(type, "__setitem__", [](VM* vm, ArgsView args){
  425. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  426. Tuple& t = CAST(Tuple&, args[1]);
  427. if(t.size() != 2){
  428. vm->TypeError("Mat3x3.__setitem__ takes a tuple of 2 integers");
  429. return vm->None;
  430. }
  431. i64 i = CAST(i64, t[0]);
  432. i64 j = CAST(i64, t[1]);
  433. if(i < 0 || i >= 3 || j < 0 || j >= 3){
  434. vm->IndexError("index out of range");
  435. return vm->None;
  436. }
  437. self.m[i][j] = CAST_F(args[2]);
  438. return vm->None;
  439. });
  440. #define PROPERTY_FIELD(field) \
  441. type->attr().set(#field, vm->property([](VM* vm, ArgsView args){ \
  442. PyMat3x3& self = _CAST(PyMat3x3&, args[0]); \
  443. return VAR(self.field); \
  444. }, [](VM* vm, ArgsView args){ \
  445. PyMat3x3& self = _CAST(PyMat3x3&, args[0]); \
  446. self.field = vm->num_to_float(args[1]); \
  447. return vm->None; \
  448. }));
  449. PROPERTY_FIELD(_11)
  450. PROPERTY_FIELD(_12)
  451. PROPERTY_FIELD(_13)
  452. PROPERTY_FIELD(_21)
  453. PROPERTY_FIELD(_22)
  454. PROPERTY_FIELD(_23)
  455. PROPERTY_FIELD(_31)
  456. PROPERTY_FIELD(_32)
  457. PROPERTY_FIELD(_33)
  458. #undef PROPERTY_FIELD
  459. vm->bind_method<1>(type, "__add__", [](VM* vm, ArgsView args){
  460. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  461. PyMat3x3& other = CAST(PyMat3x3&, args[1]);
  462. return VAR_T(PyMat3x3, self + other);
  463. });
  464. vm->bind_method<1>(type, "__sub__", [](VM* vm, ArgsView args){
  465. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  466. PyMat3x3& other = CAST(PyMat3x3&, args[1]);
  467. return VAR_T(PyMat3x3, self - other);
  468. });
  469. vm->bind_method<1>(type, "__mul__", [](VM* vm, ArgsView args){
  470. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  471. f64 other = CAST_F(args[1]);
  472. return VAR_T(PyMat3x3, self * other);
  473. });
  474. vm->bind_method<1>(type, "__truediv__", [](VM* vm, ArgsView args){
  475. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  476. f64 other = CAST_F(args[1]);
  477. return VAR_T(PyMat3x3, self / other);
  478. });
  479. auto f_mm = [](VM* vm, ArgsView args){
  480. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  481. if(is_non_tagged_type(args[1], PyMat3x3::_type(vm))){
  482. PyMat3x3& other = _CAST(PyMat3x3&, args[1]);
  483. return VAR_T(PyMat3x3, self.matmul(other));
  484. }
  485. if(is_non_tagged_type(args[1], PyVec3::_type(vm))){
  486. PyVec3& other = _CAST(PyVec3&, args[1]);
  487. return VAR_T(PyVec3, self.matmul(other));
  488. }
  489. vm->TypeError("unsupported operand type(s) for @");
  490. return vm->None;
  491. };
  492. vm->bind_method<1>(type, "__matmul__", f_mm);
  493. vm->bind_method<1>(type, "matmul", f_mm);
  494. vm->bind_method<1>(type, "__eq__", [](VM* vm, ArgsView args){
  495. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  496. PyMat3x3& other = CAST(PyMat3x3&, args[1]);
  497. return VAR(self == other);
  498. });
  499. vm->bind_method<0>(type, "determinant", [](VM* vm, ArgsView args){
  500. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  501. return VAR(self.determinant());
  502. });
  503. vm->bind_method<0>(type, "transpose", [](VM* vm, ArgsView args){
  504. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  505. return VAR_T(PyMat3x3, self.transpose());
  506. });
  507. vm->bind_method<0>(type, "inverse", [](VM* vm, ArgsView args){
  508. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  509. Mat3x3 ret;
  510. bool ok = self.inverse(ret);
  511. if(!ok) vm->ValueError("matrix is not invertible");
  512. return VAR_T(PyMat3x3, ret);
  513. });
  514. vm->bind_func<0>(type, "zeros", [](VM* vm, ArgsView args){
  515. return VAR_T(PyMat3x3, Mat3x3::zeros());
  516. });
  517. vm->bind_func<0>(type, "ones", [](VM* vm, ArgsView args){
  518. return VAR_T(PyMat3x3, Mat3x3::ones());
  519. });
  520. vm->bind_func<0>(type, "identity", [](VM* vm, ArgsView args){
  521. return VAR_T(PyMat3x3, Mat3x3::identity());
  522. });
  523. /*************** affine transformations ***************/
  524. vm->bind_func<3>(type, "trs", [](VM* vm, ArgsView args){
  525. PyVec2& t = CAST(PyVec2&, args[0]);
  526. f64 r = CAST_F(args[1]);
  527. PyVec2& s = CAST(PyVec2&, args[2]);
  528. return VAR_T(PyMat3x3, Mat3x3::trs(t, r, s));
  529. });
  530. vm->bind_method<0>(type, "is_affine", [](VM* vm, ArgsView args){
  531. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  532. return VAR(self.is_affine());
  533. });
  534. vm->bind_method<0>(type, "inverse_affine", [](VM* vm, ArgsView args){
  535. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  536. return VAR_T(PyMat3x3, self.inverse_affine());
  537. });
  538. vm->bind_method<1>(type, "matmul_affine", [](VM* vm, ArgsView args){
  539. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  540. PyMat3x3& other = CAST(PyMat3x3&, args[1]);
  541. return VAR_T(PyMat3x3, self.matmul_affine(other));
  542. });
  543. vm->bind_method<0>(type, "translation", [](VM* vm, ArgsView args){
  544. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  545. return VAR_T(PyVec2, self.translation());
  546. });
  547. vm->bind_method<0>(type, "rotation", [](VM* vm, ArgsView args){
  548. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  549. return VAR(self.rotation());
  550. });
  551. vm->bind_method<0>(type, "scale", [](VM* vm, ArgsView args){
  552. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  553. return VAR_T(PyVec2, self.scale());
  554. });
  555. vm->bind_method<1>(type, "transform_point", [](VM* vm, ArgsView args){
  556. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  557. PyVec2& v = CAST(PyVec2&, args[1]);
  558. return VAR_T(PyVec2, self.transform_point(v));
  559. });
  560. vm->bind_method<1>(type, "transform_vector", [](VM* vm, ArgsView args){
  561. PyMat3x3& self = _CAST(PyMat3x3&, args[0]);
  562. PyVec2& v = CAST(PyVec2&, args[1]);
  563. return VAR_T(PyVec2, self.transform_vector(v));
  564. });
  565. }
  566. };
  567. inline PyObject* py_var(VM* vm, Vec2 obj){ return VAR_T(PyVec2, obj); }
  568. inline PyObject* py_var(VM* vm, const PyVec2& obj){ return VAR_T(PyVec2, obj);}
  569. inline PyObject* py_var(VM* vm, Vec3 obj){ return VAR_T(PyVec3, obj); }
  570. inline PyObject* py_var(VM* vm, const PyVec3& obj){ return VAR_T(PyVec3, obj);}
  571. inline PyObject* py_var(VM* vm, const Mat3x3& obj){ return VAR_T(PyMat3x3, obj); }
  572. inline PyObject* py_var(VM* vm, const PyMat3x3& obj){ return VAR_T(PyMat3x3, obj); }
  573. inline void add_module_linalg(VM* vm){
  574. PyObject* linalg = vm->new_module("linalg");
  575. PyVec2::register_class(vm, linalg);
  576. PyVec3::register_class(vm, linalg);
  577. PyMat3x3::register_class(vm, linalg);
  578. }
  579. static_assert(sizeof(Py_<PyMat3x3>) <= 64);
  580. } // namespace pkpy